cleaned passing of kwargs, added wandb_logger to write validation metrics, log lr

Author: ayulockin

mmf/configs/defaults.yaml

@@ -54,9 +54,11 @@ training:
         # 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
+        # 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

mmf/trainers/callbacks/logistics.py

@@ -58,15 +58,12 @@ def __init__(self, config, trainer):
             if env_wandb_logdir:
                 log_dir = env_wandb_logdir
 
-            wandb_init_kwargs = config.training.wandb.init_kwargs
-
             self.wandb_logger = WandbLogger(
                 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):
@@ -157,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
@@ -400,7 +405,6 @@ class WandbLogger:
         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:
         ImportError: If wandb package is not installed.
@@ -413,7 +417,6 @@ def __init__(
         save_dir: Optional[str] = None,
         project: Optional[str] = None,
         config: Optional[Dict] = None,
-        **init_kwargs,
     ):
         try:
             import wandb
@@ -429,6 +432,11 @@ def __init__(
             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()
@@ -459,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)