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."""