log: add histogram metrics for gradients

d3rlpy/algos/qlearning/base.py

@@ -44,6 +44,7 @@
     convert_to_torch,
     convert_to_torch_recursively,
     eval_api,
+    get_gradients,
     hard_sync,
     sync_optimizer_state,
     train_api,
@@ -378,6 +379,7 @@ def fit(
         experiment_name: Optional[str] = None,
         with_timestamp: bool = True,
         logging_steps: int = 500,
+        gradient_logging_steps: Optional[int] = None,
         logging_strategy: LoggingStrategy = LoggingStrategy.EPOCH,
         logger_adapter: LoggerAdapterFactory = FileAdapterFactory(),
         show_progress: bool = True,
@@ -403,6 +405,7 @@ def fit(
                 directory name.
             logging_steps: Number of steps to log metrics. This will be ignored
                 if logging_strategy is EPOCH.
+            gradient_logging_steps: Number of steps to log gradients.
             logging_strategy: Logging strategy to use.
             logger_adapter: LoggerAdapterFactory object.
             show_progress: Flag to show progress bar for iterations.
@@ -425,6 +428,7 @@ def fit(
                 experiment_name=experiment_name,
                 with_timestamp=with_timestamp,
                 logging_steps=logging_steps,
+                gradient_logging_steps=gradient_logging_steps,
                 logging_strategy=logging_strategy,
                 logger_adapter=logger_adapter,
                 show_progress=show_progress,
@@ -442,6 +446,7 @@ def fitter(
         n_steps: int,
         n_steps_per_epoch: int = 10000,
         logging_steps: int = 500,
+        gradient_logging_steps: Optional[int] = None,
         logging_strategy: LoggingStrategy = LoggingStrategy.EPOCH,
         experiment_name: Optional[str] = None,
         with_timestamp: bool = True,
@@ -471,7 +476,8 @@ def fitter(
             with_timestamp: Flag to add timestamp string to the last of
                 directory name.
             logging_steps: Number of steps to log metrics. This will be ignored
-                if loggig_strategy is EPOCH.
+                if logging_strategy is EPOCH.
+            gradient_logging_steps: Number of steps to log gradients.
             logging_strategy: Logging strategy to use.
             logger_adapter: LoggerAdapterFactory object.
             show_progress: Flag to show progress bar for iterations.
@@ -520,6 +526,10 @@ def fitter(
         # save hyperparameters
         save_config(self, logger)
 
+        # watch model gradients
+        if gradient_logging_steps is not None:
+            logger.watch_model(gradient_logging_steps, self)
+
         # training loop
         n_epochs = n_steps // n_steps_per_epoch
         total_step = 0
@@ -559,6 +569,15 @@ def fitter(
 
                 total_step += 1
 
+                if (
+                    gradient_logging_steps is not None
+                    and total_step % gradient_logging_steps == 0
+                ):
+                    for name, grad in get_gradients(
+                        self.impl.modules.get_torch_modules()
+                    ):
+                        logger.add_histogram(name=name, values=grad)
+
                 if (
                     logging_strategy == LoggingStrategy.STEPS
                     and total_step % logging_steps == 0
@@ -608,6 +627,7 @@ def fit_online(
         experiment_name: Optional[str] = None,
         with_timestamp: bool = True,
         logging_steps: int = 500,
+        gradient_logging_steps: Optional[int] = None,
         logging_strategy: LoggingStrategy = LoggingStrategy.EPOCH,
         logger_adapter: LoggerAdapterFactory = FileAdapterFactory(),
         show_progress: bool = True,
@@ -636,6 +656,7 @@ def fit_online(
                 directory name.
             logging_steps: Number of steps to log metrics. This will be ignored
                 if logging_strategy is EPOCH.
+            gradient_logging_steps: Number of steps to log gradients.
             logging_strategy: Logging strategy to use.
             logger_adapter: LoggerAdapterFactory object.
             show_progress: Flag to show progress bar for iterations.
@@ -673,6 +694,10 @@ def fit_online(
         # save hyperparameters
         save_config(self, logger)
 
+        # watch model gradients
+        if gradient_logging_steps is not None:
+            logger.watch_model(gradient_logging_steps, self)
+
         # switch based on show_progress flag
         xrange = trange if show_progress else range
 
@@ -747,6 +772,15 @@ def fit_online(
                         ):
                             logger.commit(epoch, total_step)
 
+                if (
+                    gradient_logging_steps is not None
+                    and total_step % gradient_logging_steps == 0
+                ):
+                    for name, grad in get_gradients(
+                        self.impl.modules.get_torch_modules()
+                    ):
+                        logger.add_histogram(name=name, values=grad)
+
                 # call callback if given
                 if callback:
                     callback(self, epoch, total_step)

d3rlpy/logging/file_adapter.py

@@ -5,7 +5,14 @@
 
 import numpy as np
 
-from .logger import LOG, LoggerAdapter, LoggerAdapterFactory, SaveProtocol
+from ..types import Float32NDArray
+from .logger import (
+    LOG,
+    LoggerAdapter,
+    LoggerAdapterFactory,
+    SaveProtocol,
+    TorchModuleProtocol,
+)
 
 __all__ = ["FileAdapter", "FileAdapterFactory"]
 
@@ -60,6 +67,19 @@ def write_metric(
         with open(path, "a") as f:
             print(f"{epoch},{step},{value}", file=f)
 
+    def write_histogram(
+        self, epoch: int, step: int, name: str, value: Float32NDArray
+    ) -> None:
+        path = os.path.join(self._logdir, f"{name}.csv")
+        with open(path, "a") as f:
+            min_value = value.min()
+            max_value = value.max()
+            mean_value = value.mean()
+            print(
+                f"{epoch},{step},{name},{min_value},{max_value},{mean_value}",
+                file=f,
+            )
+
     def after_write_metric(self, epoch: int, step: int) -> None:
         pass
 
@@ -76,6 +96,11 @@ def close(self) -> None:
     def logdir(self) -> str:
         return self._logdir
 
+    def watch_model(
+        self, logging_steps: int, algo: TorchModuleProtocol
+    ) -> None:
+        pass
+
 
 class FileAdapterFactory(LoggerAdapterFactory):
     r"""FileAdapterFactory class.

d3rlpy/logging/logger.py

@@ -4,9 +4,13 @@
 from datetime import datetime
 from typing import Any, DefaultDict, Dict, Iterator, List
 
+import numpy as np
 import structlog
+from torch import nn
 from typing_extensions import Protocol
 
+from ..types import Float32NDArray
+
 __all__ = [
     "LOG",
     "set_log_context",
@@ -39,6 +43,18 @@ class SaveProtocol(Protocol):
     def save(self, fname: str) -> None: ...
 
 
+class ModuleProtocol(Protocol):
+    def get_torch_modules(self) -> List[nn.Module]: ...
+
+
+class ImplProtocol(Protocol):
+    modules: ModuleProtocol
+
+
+class TorchModuleProtocol(Protocol):
+    impl: ImplProtocol
+
+
 class LoggerAdapter(Protocol):
     r"""Interface of LoggerAdapter."""
 
@@ -69,6 +85,18 @@ def write_metric(
             value: Metric value.
         """
 
+    def write_histogram(
+        self, epoch: int, step: int, name: str, values: Float32NDArray
+    ) -> None:
+        r"""Writes histogram.
+
+        Args:
+            epoch: Epoch.
+            step: Training step.
+            name: Histogram name.
+            values: Histogram values.
+        """
+
     def after_write_metric(self, epoch: int, step: int) -> None:
         r"""Callback executed after write_metric method.
 
@@ -88,6 +116,16 @@ def save_model(self, epoch: int, algo: SaveProtocol) -> None:
     def close(self) -> None:
         r"""Closes this LoggerAdapter."""
 
+    def watch_model(
+        self, logging_steps: int, algo: TorchModuleProtocol
+    ) -> None:
+        r"""Watch model parameters / gradients during training.
+
+        Args:
+            logging_steps: Training step.
+            algo: Algorithm.
+        """
+
 
 class LoggerAdapterFactory(Protocol):
     r"""Interface of LoggerAdapterFactory."""
@@ -123,6 +161,7 @@ def __init__(
             self._experiment_name = experiment_name
         self._adapter = adapter_factory.create(self._experiment_name)
         self._metrics_buffer = defaultdict(list)
+        self._histogram_metrics_buffer = defaultdict(list)
 
     def add_params(self, params: Dict[str, Any]) -> None:
         self._adapter.write_params(params)
@@ -131,6 +170,9 @@ def add_params(self, params: Dict[str, Any]) -> None:
     def add_metric(self, name: str, value: float) -> None:
         self._metrics_buffer[name].append(value)
 
+    def add_histogram(self, name: str, values: Float32NDArray) -> None:
+        self._histogram_metrics_buffer[name].append(values)
+
     def commit(self, epoch: int, step: int) -> Dict[str, float]:
         self._adapter.before_write_metric(epoch, step)
 
@@ -140,6 +182,10 @@ def commit(self, epoch: int, step: int) -> Dict[str, float]:
             self._adapter.write_metric(epoch, step, name, metric)
             metrics[name] = metric
 
+        for name, buffer in self._histogram_metrics_buffer.items():
+            histogram_values = np.concatenate(buffer)
+            self._adapter.write_histogram(epoch, step, name, histogram_values)
+
         LOG.info(
             f"{self._experiment_name}: epoch={epoch} step={step}",
             epoch=epoch,
@@ -171,3 +217,6 @@ def measure_time(self, name: str) -> Iterator[None]:
     @property
     def adapter(self) -> LoggerAdapter:
         return self._adapter
+
+    def watch_model(self, logging_steps: int, algo: TorchModuleProtocol) -> None:
+        self._adapter.watch_model(logging_steps, algo)

d3rlpy/logging/noop_adapter.py

@@ -1,6 +1,12 @@
 from typing import Any, Dict
 
-from .logger import LoggerAdapter, LoggerAdapterFactory, SaveProtocol
+from ..types import Float32NDArray
+from .logger import (
+    LoggerAdapter,
+    LoggerAdapterFactory,
+    SaveProtocol,
+    TorchModuleProtocol,
+)
 
 __all__ = ["NoopAdapter", "NoopAdapterFactory"]
 
@@ -23,6 +29,11 @@ def write_metric(
     ) -> None:
         pass
 
+    def write_histogram(
+        self, epoch: int, step: int, name: str, values: Float32NDArray
+    ) -> None:
+        pass
+
     def after_write_metric(self, epoch: int, step: int) -> None:
         pass
 
@@ -32,6 +43,11 @@ def save_model(self, epoch: int, algo: SaveProtocol) -> None:
     def close(self) -> None:
         pass
 
+    def watch_model(
+        self, logging_steps: int, algo: TorchModuleProtocol
+    ) -> None:
+        pass
+
 
 class NoopAdapterFactory(LoggerAdapterFactory):
     r"""NoopAdapterFactory class.

d3rlpy/logging/tensorboard_adapter.py

@@ -3,7 +3,13 @@
 
 import numpy as np
 
-from .logger import LoggerAdapter, LoggerAdapterFactory, SaveProtocol
+from ..types import Float32NDArray
+from .logger import (
+    LoggerAdapter,
+    LoggerAdapterFactory,
+    SaveProtocol,
+    TorchModuleProtocol,
+)
 
 __all__ = ["TensorboardAdapter", "TensorboardAdapterFactory"]
 
@@ -50,6 +56,11 @@ def write_metric(
         self._writer.add_scalar(f"metrics/{name}", value, epoch)
         self._metrics[name] = value
 
+    def write_histogram(
+        self, epoch: int, step: int, name: str, value: Float32NDArray
+    ) -> None:
+        self._writer.add_histogram(f"histograms/{name}_grad", value, epoch)
+
     def after_write_metric(self, epoch: int, step: int) -> None:
         self._writer.add_hparams(
             self._params,
@@ -64,6 +75,11 @@ def save_model(self, epoch: int, algo: SaveProtocol) -> None:
     def close(self) -> None:
         self._writer.close()
 
+    def watch_model(
+        self, logging_steps: int, algo: TorchModuleProtocol
+    ) -> None:
+        pass
+
 
 class TensorboardAdapterFactory(LoggerAdapterFactory):
     r"""TensorboardAdapterFactory class.

d3rlpy/logging/utils.py

@@ -1,6 +1,12 @@
 from typing import Any, Dict, Sequence
 
-from .logger import LoggerAdapter, LoggerAdapterFactory, SaveProtocol
+from ..types import Float32NDArray
+from .logger import (
+    LoggerAdapter,
+    LoggerAdapterFactory,
+    SaveProtocol,
+    TorchModuleProtocol,
+)
 
 __all__ = ["CombineAdapter", "CombineAdapterFactory"]
 
@@ -32,6 +38,12 @@ def write_metric(
         for adapter in self._adapters:
             adapter.write_metric(epoch, step, name, value)
 
+    def write_histogram(
+        self, epoch: int, step: int, name: str, values: Float32NDArray
+    ) -> None:
+        for adapter in self._adapters:
+            adapter.write_histogram(epoch, step, name, values)
+
     def after_write_metric(self, epoch: int, step: int) -> None:
         for adapter in self._adapters:
             adapter.after_write_metric(epoch, step)
@@ -44,6 +56,12 @@ def close(self) -> None:
         for adapter in self._adapters:
             adapter.close()
 
+    def watch_model(
+        self, logging_steps: int, algo: TorchModuleProtocol
+    ) -> None:
+        for adapter in self._adapters:
+            adapter.watch_model(logging_steps, algo)
+
 
 class CombineAdapterFactory(LoggerAdapterFactory):
     r"""CombineAdapterFactory class.