Feature: wandb logger

.run_local_tests.sh

@@ -3,6 +3,7 @@ rm -rf _ckpt_*
 rm -rf tests/save_dir*
 rm -rf tests/mlruns_*
 rm -rf tests/cometruns*
+rm -rf tests/wandb*
 rm -rf tests/tests/*
 rm -rf lightning_logs
 coverage run --source pytorch_lightning -m py.test pytorch_lightning tests pl_examples -v --doctest-modules

README.md

@@ -306,6 +306,7 @@ Lightning also adds a text column with all the hyperparameters for this experime
 - [Save a snapshot of all hyperparameters](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#save-a-snapshot-of-all-hyperparameters) 
 - [Snapshot code for a training run](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#snapshot-code-for-a-training-run) 
 - [Write logs file to csv every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#write-logs-file-to-csv-every-k-batches)
+- [Logging on W&B](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#w&b)
 
 #### Training loop 
 

docs/source/conf.py

@@ -302,7 +302,7 @@ def setup(app):
  MOCK_REQUIRE_PACKAGES.append(pkg.rstrip())
 
 # TODO: better parse from package since the import name and package name may differ
-MOCK_MANUAL_PACKAGES = ['torch', 'torchvision', 'sklearn', 'test_tube', 'mlflow', 'comet_ml']
+MOCK_MANUAL_PACKAGES = ['torch', 'torchvision', 'sklearn', 'test_tube', 'mlflow', 'comet_ml', 'wandb']
 autodoc_mock_imports = MOCK_REQUIRE_PACKAGES + MOCK_MANUAL_PACKAGES
 # for mod_name in MOCK_REQUIRE_PACKAGES:
 # sys.modules[mod_name] = mock.Mock()

pytorch_lightning/logging/__init__.py

@@ -180,6 +180,10 @@ def __init__(self, hparams):
 except ImportError:
  pass
 
+try:
+ from .wandb import WandbLogger
+except ImportError:
+ pass
 try:
  # needed to prevent ImportError and duplicated logs.
  environ["COMET_DISABLE_AUTO_LOGGING"] = "1"

pytorch_lightning/logging/wandb.py

@@ -0,0 +1,116 @@
+"""
+Log using `W&B <https://www.wandb.com>`_
+
+.. code-block:: python
+
+ from pytorch_lightning.logging import WandbLogger
+ wandb_logger = WandbLogger(
+ name="my_new_nun", # Optional, display name
+ save_dir="my_folder", # Optional, path to save data
+ offline=False, # Optional, run offline (can sync later)
+ version="run_12345", # Optional, used to restart previous run
+ id="run_12345", # Optional, same as version
+ anonymous=False, # Optional, enable or disable anonymous logging
+ project="bert", # Optional, project to which run belongs to
+ tags=["tag1", "tag2"] # Optional, tags associated with run
+ )
+ trainer = Trainer(logger=wandb_logger)
+
+
+Use the logger anywhere in you LightningModule as follows:
+
+.. code-block:: python
+
+ def train_step(...):
+ # example
+ self.logger.experiment.whatever_wandb_supports(...)
+
+ def any_lightning_module_function_or_hook(...):
+ self.logger.experiment.whatever_wandb_supports(...)
+
+"""
+
+import os
+
+try:
+ import wandb
+except ImportError:
+ raise ImportError('Missing wandb package.')
+
+from .base import LightningLoggerBase, rank_zero_only
+
+
+class WandbLogger(LightningLoggerBase):
+ """
+ Logger for W&B.
+
+ Args:
+ name (str): display name for the run.
+ save_dir (str): path where data is saved.
+ offline (bool): run offline (data can be streamed later to wandb servers).
+ id or version (str): sets the version, mainly used to resume a previous run.
+ anonymous (bool): enables or explicitly disables anonymous logging.
+ project (str): the name of the project to which this run will belong.
+ tags (list of str): tags associated with this run.
+ """
+
+ def __init__(self, name=None, save_dir=None, offline=False, id=None, anonymous=False,
+ version=None, project=None, tags=None):
+ super().__init__()
+ self._name = name
+ self._save_dir = save_dir
+ self._anonymous = "allow" if anonymous else None
+ self._id = version or id
+ self._tags = tags
+ self._project = project
+ self._experiment = None
+ self._offline = offline
+
+ def __getstate__(self):
+ state = self.__dict__.copy()
+ # cannot be pickled
+ state['_experiment'] = None
+ # args needed to reload correct experiment
+ state['_id'] = self.experiment.id
+ return state
+
+ @property
+ def experiment(self):
+ if self._experiment is None:
+ if self._offline:
+ os.environ["WANDB_MODE"] = "dryrun"
+ self._experiment = wandb.init(
+ name=self._name, dir=self._save_dir, project=self._project, anonymous=self._anonymous,
+ id=self._id, resume="allow", tags=self._tags)
+ return self._experiment
+
+ def watch(self, model, log="gradients", log_freq=100):
+ wandb.watch(model, log, log_freq)
+
+ @rank_zero_only
+ def log_hyperparams(self, params):
+ self.experiment.config.update(params)
+
+ @rank_zero_only
+ def log_metrics(self, metrics, step=None):
+ metrics["global_step"] = step
+ self.experiment.history.add(metrics)
+
+ def save(self):
+ pass
+
+ @rank_zero_only
+ def finalize(self, status='success'):
+ try:
+ exit_code = 0 if status == 'success' else 1
+ wandb.join(exit_code)
+ except TypeError:
+ wandb.join()
+
+ @property
+ def name(self):
+ return self.experiment.project_name()
+
+ @property
+ def version(self):
+ return self.experiment.id

tests/requirements.txt

@@ -8,4 +8,5 @@ check-manifest
 # test_tube # already installed in main req.
 mlflow
 comet_ml
+wandb
 twine==1.13.0

tests/test_logging.py

@@ -193,6 +193,63 @@ def test_comet_pickle(tmpdir, monkeypatch):
  trainer2.logger.log_metrics({"acc": 1.0})
 
 
+def test_wandb_logger(tmpdir):
+ """Verify that basic functionality of wandb logger works."""
+ tutils.reset_seed()
+
+ try:
+ from pytorch_lightning.logging import WandbLogger
+ except ModuleNotFoundError:
+ return
+
+ hparams = tutils.get_hparams()
+ model = LightningTestModel(hparams)
+
+ wandb_dir = os.path.join(tmpdir, "wandb")
+
+ logger = WandbLogger(save_dir=wandb_dir, anonymous=True)
+
+ trainer_options = dict(
+ default_save_path=tmpdir,
+ max_epochs=1,
+ train_percent_check=0.01,
+ logger=logger
+ )
+ trainer = Trainer(**trainer_options)
+ result = trainer.fit(model)
+
+ print('result finished')
+ assert result == 1, "Training failed"
+
+
+def test_wandb_pickle(tmpdir):
+ """Verify that pickling trainer with wandb logger works."""
+ tutils.reset_seed()
+
+ try:
+ from pytorch_lightning.logging import WandbLogger
+ except ModuleNotFoundError:
+ return
+
+ # hparams = tutils.get_hparams()
+ # model = LightningTestModel(hparams)
+
+ wandb_dir = os.path.join(tmpdir, "wandb")
+
+ logger = WandbLogger(save_dir=wandb_dir, anonymous=True)
+
+ trainer_options = dict(
+ default_save_path=tmpdir,
+ max_epochs=1,
+ logger=logger
+ )
+
+ trainer = Trainer(**trainer_options)
+ pkl_bytes = pickle.dumps(trainer)
+ trainer2 = pickle.loads(pkl_bytes)
+ trainer2.logger.log_metrics({"acc": 1.0})
+
+
 def test_tensorboard_logger(tmpdir):
  """Verify that basic functionality of Tensorboard logger works."""