Add logger to Weights & Biases

CHANGES.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added `NeptuneLogger` callback for logging experiment metadata to neptune.ai
 - Add DataFrameTransformer, an sklearn compatible transformer that helps working with pandas DataFrames by transforming the DataFrame into a representation that works well with neural networks (#507)
+- Added `WandbLogger` callback for logging to Weights & Biases
 
 ### Changed
 

skorch/callbacks/__init__.py

@@ -33,5 +33,6 @@
     'TrainEndCheckpoint',
     'TensorBoard',
     'Unfreezer',
+    'WandbLogger',
     'WarmRestartLR',
 ]

skorch/callbacks/logging.py

@@ -5,6 +5,7 @@
 from contextlib import suppress
 from numbers import Number
 from itertools import cycle
+from pathlib import Path
 
 import numpy as np
 import tqdm
@@ -14,7 +15,8 @@
 from skorch.dataset import get_len
 from skorch.callbacks import Callback
 
-__all__ = ['EpochTimer', 'NeptuneLogger', 'PrintLog', 'ProgressBar', 'TensorBoard']
+__all__ = ['EpochTimer', 'NeptuneLogger', 'WandbLogger', 'PrintLog', 'ProgressBar',
+           'TensorBoard']
 
 
 def filter_log_keys(keys, keys_ignored=None):
@@ -205,6 +207,76 @@ def on_train_end(self, net, **kwargs):
         if self.close_after_train:
             self.experiment.stop()
 
+class WandbLogger(Callback):
+    """Logs best model and metrics to `Weights & Biases <https://docs.wandb.com/>`_
+
+    "Use this callback to automatically log best trained model and all metrics from
+    your net's history to Weights & Biases after each epoch.
+
+    Examples
+    --------
+    >>> import wandb
+    >>> from skorch.callbacks import WandbLogger
+    >>> wandb_run = wandb.init()
+    >>> wandb.config.update({"learning rate": 1e-3, "batch size": 32})  # optional
+    >>> net = NeuralNet(..., callbacks=[WandbLogger(wandb_run)])
+    >>> net.fit(X, y)
+
+    Parameters
+    ----------
+    wandb_run : wandb.wandb_run.Run
+      wandb Run used to log data.
+
+    save_model : bool (default=True)
+      Saves best trained model.
+
+    keys_ignored : str or list of str (default=None)
+      Key or list of keys that should not be logged to
+      tensorboard. Note that in addition to the keys provided by the
+      user, keys such as those starting with 'event_' or ending on
+      '_best' are ignored by default.
+    """
+
+    # Record if watch has been called previously (even in another instance)
+    _watch_called = False
+
+    def __init__(
+            self,
+            wandb_run,
+            save_model=True,
+            keys_ignored=None,
+    ):
+        self.wandb_run = wandb_run
+        self.save_model = save_model
+        self.keys_ignored = keys_ignored
+        self.model_path = Path(wandb_run.dir) / 'best_model.pth'
+
+    def initialize(self):
+        keys_ignored = self.keys_ignored
+        if isinstance(keys_ignored, str):
+            keys_ignored = [keys_ignored]
+        self.keys_ignored_ = set(keys_ignored or [])
+        self.keys_ignored_.add('batches')
+        return self
+
+    def on_train_begin(self, net, **kwargs):
+        """Log model topology and add a hook for gradients"""
+        if not WandbLogger._watch_called:
+            WandbLogger._watch_called = True
+            self.wandb_run.watch(net.module_)
+
+    def on_epoch_end(self, net, **kwargs):
+        """Automatically log values from the last history step."""
+        hist = net.history[-1]
+        keys_kept = filter_log_keys(hist, keys_ignored=self.keys_ignored_)
+        logged_vals = dict((k, hist[k]) for k in keys_kept if k in hist)
+        self.wandb_run.log(logged_vals)
+
+        # save best model
+        if self.save_model and hist['valid_loss_best']:
+            with self.model_path.open('wb') as model_file:
+                net.save_params(f_params=model_file)
+
 
 class PrintLog(Callback):
     """Print useful information from the model's history as a table.