[refactor] Extend WandbLogger to log config variables, entity and kwargs

mmf/configs/defaults.yaml

@@ -45,11 +45,21 @@ training:
     wandb:
         # Whether to use Weights and Biases Logger, (Default: false)
         enabled: false
+        # An entity is a username or team name where you're sending runs.
+        # This is necessary if you want to log your metrics to a team account. By default
+        # it will log the run to your user account.
+        entity: null
         # Project name to be used while logging the experiment with wandb
         wandb_projectname: mmf_${oc.env:USER,}
         # Experiment/ run name to be used while logging the experiment
         # under the project with wandb
         wandb_runname: ${training.experiment_name}
+        # Specify other argument values that you want to pass to wandb.init(). Check out the documentation
+        # at https://docs.wandb.ai/ref/python/init to see what arguments are available.
+        # job_type: 'train'
+        # tags: ['tag1', 'tag2']
+
+
 
     # Size of the batch globally. If distributed or data_parallel
     # is used, this will be divided equally among GPUs

mmf/trainers/callbacks/logistics.py

@@ -58,11 +58,12 @@ def __init__(self, config, trainer):
             if env_wandb_logdir:
                 log_dir = env_wandb_logdir
 
-            wandb_projectname = config.training.wandb.wandb_projectname
-            wandb_runname = config.training.wandb.wandb_runname
-
             self.wandb_logger = WandbLogger(
-                name=wandb_runname, save_dir=log_dir, project=wandb_projectname
+                entity=config.training.wandb.entity,
+                project=config.training.wandb.wandb_projectname,
+                config=config,
+                name=config.training.wandb.wandb_runname,
+                save_dir=log_dir,
             )
 
     def on_train_start(self):
@@ -153,6 +154,7 @@ def on_test_end(self, **kwargs):
             meter=kwargs["meter"],
             should_print=prefix,
             tb_writer=self.tb_writer,
+            wandb_logger=self.wandb_logger,
         )
         logger.info(f"Finished run in {self.total_timer.get_time_since_start()}")
 

mmf/utils/logger.py

@@ -2,6 +2,7 @@
 
 import collections
 import functools
+import itertools
 import json
 import logging
 import os
@@ -231,7 +232,11 @@ def summarize_report(
 
     if wandb_logger:
         metrics = meter.get_scalar_dict()
-        wandb_logger.log_metrics({**metrics, "trainer/global_step": current_iteration})
+        wandb_logger.log_metrics({**metrics, "trainer/global_step": current_iteration}, commit=False)
+
+    # Log the learning rate if available
+    if wandb_logger and 'lr' in extra.keys():
+        wandb_logger.log_metrics({"train/learning_rate": float(extra["lr"])})
 
     if not should_print:
         return
@@ -395,21 +400,23 @@ class WandbLogger:
     Log using `Weights and Biases`.
 
     Args:
+        entity: An entity is a username or team name where you're sending runs.
         name: Display name for the run.
         save_dir: Path where data is saved (./save/logs/wandb/ by default).
         project: Display name for the project.
-        **init_kwargs: Arguments passed to :func:`wandb.init`.
+        config: Configuration for the run.
 
     Raises:
         ImportError: If wandb package is not installed.
     """
 
     def __init__(
         self,
+        entity: Optional[str] = None,
         name: Optional[str] = None,
         save_dir: Optional[str] = None,
         project: Optional[str] = None,
-        **init_kwargs,
+        config: Optional[Dict] = None,
     ):
         try:
             import wandb
@@ -421,8 +428,15 @@ def __init__(
 
         self._wandb = wandb
 
-        self._wandb_init = dict(name=name, project=project, dir=save_dir)
+        self._wandb_init = dict(
+            entity=entity, name=name, project=project, dir=save_dir, config=config
+        )
 
+        init_kwargs = dict(
+            itertools.islice(
+                config.training.wandb.items(), 4, len(config.training.wandb)
+            )
+        )
         self._wandb_init.update(**init_kwargs)
 
         self.setup()
@@ -453,14 +467,15 @@ def _should_log_wandb(self):
         else:
             return True
 
-    def log_metrics(self, metrics: Dict[str, float]):
+    def log_metrics(self, metrics: Dict[str, float], commit=True):
         """
         Log the monitored metrics to the wand dashboard.
 
         Args:
-            metrics (Dict[str, float]): [description]
+            metrics (Dict[str, float]): A dictionary of metrics to log.
+            commit (bool): Save the metrics dict to the wandb server and increment the step. (default: True)
         """
         if not self._should_log_wandb():
             return
 
-        self._wandb.log(metrics)
+        self._wandb.log(metrics, commit=commit)

website/docs/notes/logging.md

@@ -1,42 +1,75 @@
 ---
-id: concepts
-title: Terminology and Concepts
-sidebar_label: Terminology and Concepts
+id: logger
+title: Weights and Biases Logging
+sidebar_label: Weights and Biases Logging
 ---
 
 ## Weights and Biases Logger
 
-MMF has a `WandbLogger` class which lets the user to log their model's progress using [Weights and Biases](https://gitbook-docs.wandb.ai/).
+MMF now has a `WandbLogger` class which lets the user to log their model's progress using [Weights and Biases](https://wandb.ai/site). Enable this logger to automatically log the training/validation metrics, system (GPU and CPU) metrics and configuration parameters.
+
+## First time setup 
 
 To set up wandb, run the following:
 ```
 pip install wandb
+```
+In order to log anything to the W&B server you need to authenticate the machine with W&B **API key**. You can create a new account by going to https://wandb.ai/signup which will generate an API key. If you are an existing user you can retrieve your key from https://wandb.ai/authorize. You only need to supply your key once, and then it is remembered on the same device. 
+
+```
 wandb login
 ```
 
+## W&B config parameters
+
 The following options are available in config to enable and customize the wandb logging:
 ```yaml
 training:
     # Weights and Biases control, by default Weights and Biases (wandb) is disabled
     wandb:
         # Whether to use Weights and Biases Logger, (Default: false)
         enabled: false
+        # An entity is a username or team name where you're sending runs.
+        # This is necessary if you want to log your metrics to a team account. By default
+        # it will log the run to your user account.
+        entity: null
         # Project name to be used while logging the experiment with wandb
-        wandb_projectname: mmf_${oc.env:USER}
+        wandb_projectname: mmf_${oc.env:USER,}
         # Experiment/ run name to be used while logging the experiment
         # under the project with wandb
         wandb_runname: ${training.experiment_name}
+        # Specify other argument values that you want to pass to wandb.init(). Check out the documentation
+        # at https://docs.wandb.ai/ref/python/init to see what arguments are available.
+        # job_type: 'train'
+        # tags: ['tag1', 'tag2']
 env:
     wandb_logdir: ${env:MMF_WANDB_LOGDIR,}
-```
-To enable wandb logger the user needs to change the following option in the config.
+``` 
+
+* To enable wandb logger the user needs to change the following option in the config.
+
+	`training.wandb.enabled=True`
+
+* To give the `entity` which is the name of the team or the username, the user needs to change the following option in the config. In case no `entity` is provided, the data will be logged to the `entity` set as default in the user's settings. 
+
+	`training.wandb.entity=<teamname/username>`
+
+* To give the current experiment a project and run name, user should add these config options.
+
+	`training.wandb.wandb_projectname=<ProjectName>` <br>
+	`training.wandb.wandb_runname=<RunName>`
+
+* To change the path to the directory where wandb metadata would be stored (Default: `env.log_dir`):
 
-`training.wandb.enabled=True`
+	`env.wandb_logdir=<dir_name>`
 
-To give the current experiment a project and run name, user should add these config options.
+* To provide extra arguments to `wandb.init()`, the user just needs to define them in the config file. Check out the documentation at https://docs.wandb.ai/ref/python/init to see what arguments are available. An example is shown in the config parameter shown above. 
 
-`training.wandb.wandb_projectname=<ProjectName> training.wandb.wandb_runname=<RunName>`
+## Current features
 
-To change the path to the directory where wandb metadata would be stored (Default: `env.log_dir`):
+The following features are currently supported by the `WandbLogger`:
 
-`env.wandb_logdir=<dir_name>`
+* Training & Validation metrics
+* Learning Rate over time
+* GPU: Type, GPU Utilization, power, temperature, CUDA memory usage
+* Log configuration parameters