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

mmf/configs/defaults.yaml

@@ -45,11 +45,19 @@ 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 to be used while logging the experiment
+        init_kwargs:
+            job_type: train
+
 
     # 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,15 @@ 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
+            wandb_init_kwargs = config.training.wandb.init_kwargs
 
             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,
+                **wandb_init_kwargs,
             )
 
     def on_train_start(self):

mmf/utils/logger.py

@@ -395,9 +395,11 @@ 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.
+        config: Configuration for the run.
         **init_kwargs: Arguments passed to :func:`wandb.init`.
 
     Raises:
@@ -406,9 +408,11 @@ class WandbLogger:
 
     def __init__(
         self,
+        entity: Optional[str] = None,
         name: Optional[str] = None,
         save_dir: Optional[str] = None,
         project: Optional[str] = None,
+        config: Optional[Dict] = None,
         **init_kwargs,
     ):
         try:
@@ -421,7 +425,9 @@ 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
+        )
 
         self._wandb_init.update(**init_kwargs)