add documentation for W&B logger (#74)

Co-authored-by: Zhaocheng Zhu <healmysoul@163.com>

Author: manangoel99

README.md

@@ -78,7 +78,8 @@ Quick Start
 
 TorchDrug is designed for humans and focused on graph structured data.
 It enables easy implementation of graph operations in machine learning models.
-All the operations in TorchDrug are backed by [PyTorch] framework, and support GPU acceleration and auto differentiation.
+All the operations in TorchDrug are backed by [PyTorch] framework, and support GPU
+acceleration and auto differentiation.
 
 ```python
 from torchdrug import data
@@ -90,7 +91,8 @@ graph = graph.cuda()
 subgraph = graph.subgraph([2, 3, 4])
 ```
 
-Molecules are also supported in TorchDrug. You can get the desired molecule properties without any domain knowledge.
+Molecules are also supported in TorchDrug. You can get the desired molecule
+properties without any domain knowledge.
 
 ```python
 mol = data.Molecule.from_smiles("CCOC(=O)N", node_feature="default", edge_feature="default")
@@ -99,7 +101,8 @@ print(mol.atom_type)
 print(mol.to_scaffold())
 ```
 
-You may also register custom node, edge or graph attributes. They will be automatically processed during indexing operations.
+You may also register custom node, edge or graph attributes. They will be
+automatically processed during indexing operations.
 
 ```python
 with mol.edge():
@@ -108,8 +111,9 @@ sub_mol = mol.subgraph(mol.atom_type != td.NITROGEN)
 print(sub_mol.is_CC_bond)
 ```
 
-TorchDrug provides a wide range of common datasets and building blocks for drug discovery.
-With minimal code, you can apply standard models to solve your own problem.
+TorchDrug provides a wide range of common datasets and building blocks for drug
+discovery. With minimal code, you can apply standard models to solve your own
+problem.
 
 ```python
 import torch
@@ -144,6 +148,13 @@ solver = core.Engine(task, train_set, valid_set, test_set, optimizer, gpus=[0, 1
 solver = core.Engine(task, train_set, valid_set, test_set, optimizer, gpus=[0, 1, 2, 3, 0, 1, 2, 3])
 ```
 
+Experiments can be easily tracked and managed through [Weights & Biases platform].
+```python
+solver = core.Engine(task, train_set, valid_set, test_set, optimizer, logger="wandb")
+```
+
+[Weights & Biases platform]: https://wandb.ai/
+
 Contributing
 ------------
 

doc/source/quick_start.rst

@@ -177,6 +177,14 @@ the script using ``python -m torch.distributed.launch --nproc_per_node=4``.
     solver = core.Engine(task, train_set, valid_set, test_set, optimizer,
                          batch_size=256, gpus=[0, 1, 2, 3])
 
+We may log the training and evaluation metrics to Weights & Biases platform for
+better experiment tracking in the browser.
+
+.. code:: python
+
+    solver = core.Engine(task, train_set, valid_set, test_set, optimizer,
+                         batch_size=1024, logger="wandb")
+
 Test the Model
 ^^^^^^^^^^^^^^
 

torchdrug/core/__init__.py

@@ -1,9 +1,9 @@
 from .core import _MetaContainer, Registry, Configurable, make_configurable
 from .engine import Engine
 from .meter import Meter
-from .logger import LoggerBase, ConsoleLogger, WandbLogger
+from .logger import LoggerBase, LoggingLogger, WandbLogger
 
 __all__ = [
     "_MetaContainer", "Registry", "Configurable",
-    "Engine", "Meter", "LoggerBase", "ConsoleLogger", "WandbLogger",
+    "Engine", "Meter", "LoggerBase", "LoggingLogger", "WandbLogger",
 ]

torchdrug/core/engine.py

@@ -53,12 +53,13 @@ class Engine(core.Configurable):
         gradient_interval (int, optional): perform a gradient update every n batches.
             This creates an equivalent batch size of ``batch_size * gradient_interval`` for optimization.
         num_worker (int, optional): number of CPU workers per GPU
-        logger (str or core.LoggerBase, optional): logger type. Avaiable types are ``console`` and ``wandb``.
+        logger (str or core.LoggerBase, optional): logger type or logger instance.
+            Available types are ``logging`` and ``wandb``.
         log_interval (int, optional): log every n gradient updates
     """
 
     def __init__(self, task, train_set, valid_set, test_set, optimizer, scheduler=None, gpus=None, batch_size=1,
-                 gradient_interval=1, num_worker=0, logger="console", log_interval=100):
+                 gradient_interval=1, num_worker=0, logger="logging", log_interval=100):
         self.rank = comm.get_rank()
         self.world_size = comm.get_world_size()
         self.gpus = gpus
@@ -107,10 +108,12 @@ def __init__(self, task, train_set, valid_set, test_set, optimizer, scheduler=No
         self.scheduler = scheduler
 
         if isinstance(logger, str):
-            if logger == "console":
-                logger = core.ConsoleLogger()
+            if logger == "logging":
+                logger = core.LoggingLogger()
             elif logger == "wandb":
                 logger = core.WandbLogger(project=task.__class__.__name__)
+            else:
+                raise ValueError("Unknown logger `%s`" % logger)
         self.meter = core.Meter(log_interval=log_interval, silent=self.rank > 0, logger=logger)
         self.meter.log_config(self.config_dict())
 

torchdrug/core/logger.py

@@ -1,5 +1,6 @@
 import pprint
 import logging
+import warnings
 
 from torchdrug.core import Registry as R
 from torchdrug.utils import pretty
@@ -34,10 +35,23 @@ def log_config(self, config):
         raise NotImplementedError
 
 
-@R.register("core.ConsoleLogger")
-class ConsoleLogger(LoggerBase):
+@R.register("core.LoggingLogger")
+class LoggingLogger(LoggerBase):
     """
-    Logger for console output.
+    Log outputs with the builtin logging module of Python.
+
+    By default, the logs will be printed to the console. To additionally log outputs to a file,
+    add the following lines in the beginning of your code.
+
+    .. code-block: python
+
+        import logging
+
+        format = logging.Formatter("%(asctime)-10s %(message)s", "%H:%M:%S")
+        handler = logging.FileHandler("log.txt")
+        handler.setFormatter(format)
+        logger = logging.getLogger("")
+        logger.addHandler(handler)
     """
 
     def __init__(self):
@@ -60,14 +74,25 @@ def log_config(self, config):
 
 
 @R.register("core.WandbLogger")
-class WandbLogger(ConsoleLogger):
+class WandbLogger(LoggingLogger):
     """
-    Logger for wandb and console outputs.
+    Log outputs with `Weights and Biases`_ and track the experiment progress.
+
+    Note this class also output logs with the builtin logging module.
+
+    See `wandb.init`_ for more details.
+
+    .. _Weights and Biases:
+        https://docs.wandb.ai
+
+    .. _wandb.init:
+        https://docs.wandb.ai/ref/python/init
 
     Parameters:
-        project (str, optional): name of the project in wandb
-        name (str, optional): name for this run in wandb
-        dir (str, optional): path to save wandb outputs. By default, outputs are stored in ``./wandb``.
+        project (str, optional): name of the project
+        name (str, optional): name of this run
+        dir (str, optional): path to store meta data. Default is `./wandb`.
+        kwargs: keyword arguments for `wandb.init`_
     """
 
     def __init__(self, project=None, name=None, dir=None, **kwargs):
@@ -77,7 +102,14 @@ def __init__(self, project=None, name=None, dir=None, **kwargs):
         except ModuleNotFoundError:
             raise ModuleNotFoundError("Wandb is not found. Please install it with `pip install wandb`")
 
-        self.run = wandb.init(project=project, name=name, dir=dir, reinit=True, **kwargs)
+        if wandb.run is not None:
+            warnings.warn(
+                 "There is a wandb run already in progress and newly created instances of `WandbLogger` will reuse"
+                " this run. If this is not desired, call `wandb.finish()` or `WandbLogger.finish()` before instantiating `WandbLogger`."
+            )
+            self.run = wandb.run
+        else:
+            self.run = wandb.init(project=project, name=name, dir=dir, **kwargs)
 
         self.run.define_metric("train/batch/*", step_metric="batch", summary="none")
         for split in ["train", "valid", "test"]: