[feat] Model version control using W&B Artifacts (#1137)

Summary:
🚀 I have extended the `WandbLogger` with the ability to log the `current.pt` checkpoint as W&B Artifacts. Note that this PR is based on top of this [PR](#1129).

### What is W&B Artifacts?

> W&B Artifacts was designed to make it effortless to version your datasets and models, regardless of whether you want to store your files with us or whether you already have a bucket you want us to track. Once you've tracked your dataset or model files, W&B will automatically log each and every modification, giving you a complete and auditable history of changes to your files.

Through this PR, W&B Artifacts can help save and organize machine learning models throughout a project's lifecycle. More details in the documentation [here](https://docs.wandb.ai/guides/artifacts/model-versioning).

### Modification

This PR adds a `log_model_checkpoint` method to the `WandbLogger` class in the `utils/logger.py` file. This method is called in the `utils/checkpoint.py` file.

### Usage

To use this, in the `config/defaults.yaml` do, `training.wandb.enabled=true` and `training.wandb.log_checkpoint=true`.

### Result

The screenshot shows the `current.pt` checkpoints saved at intervals defined by `training.checkpoint_interval`. You can check out the logged artifacts page [here](https://wandb.ai/ayut/mmf/artifacts/model/run_ey9xextf_model/0dc64164acbdc300fd01/api).

![image](https://user-images.githubusercontent.com/31141479/139390462-d5c8445e-5c20-4fdd-85d0-51ef64846bf0.png)

### Superpowers

With this small addition, now one can easily track different versions of the model, download a checkpoint of interest by using the API in the API tab, easily share the checkpoints with teammates, etc.

### Requests

This is a draft PR as there are a few more things that can be improved here.

* Is there a better way to access the path to the `current.pt` checkpoint? Rather is the modification made to `utils/checkpoint.py` an acceptable way of approaching this?

* While logging a file as W&B artifacts we can also provide metadata associated with that file. In this case, we can add current iteration, training metrics, etc. as the metadata. Would love to get suggestions about the different data points that I should log as metadata alongside the checkpoints.

* How to determine if a checkpoint is the best one? If a checkpoint is best I can add `best` as an alias for that checkpoint's artifact.

Pull Request resolved: #1137

Test Plan:
Imported from GitHub, without a `Test Plan:` line.

**Static Docs Preview: mmf**
|[Full Site](https://our.intern.facebook.com/intern/staticdocs/eph/D32402090/V6/mmf/)|

|**Modified Pages**|
|[docs/notes/logger](https://our.intern.facebook.com/intern/staticdocs/eph/D32402090/V6/mmf/docs/notes/logger/)|

Reviewed By: apsdehal

Differential Revision: D32402090

Pulled By: ebsmothers

fbshipit-source-id: 94b881ec55c4197301331d571bc926521e2feecc

Author: ayulockin

mmf/configs/defaults.yaml

@@ -45,11 +45,24 @@ 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,}
+        project: mmf
         # Experiment/ run name to be used while logging the experiment
         # under the project with wandb
-        wandb_runname: ${training.experiment_name}
+        name: ${training.experiment_name}
+        # You can save your model checkpoints as W&B Artifacts for model versioning.
+        # Set the value to `true` to enable this feature.
+        log_checkpoint: false
+        # 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,10 @@ 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,
+                config=config,
+                project=config.training.wandb.project,
             )
 
     def on_train_start(self):
@@ -153,6 +152,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/checkpoint.py

@@ -522,6 +522,7 @@ def save(self, update, iteration=None, update_best=False):
         best_metric = (
             self.trainer.early_stop_callback.early_stopping.best_monitored_value
         )
+
         model = self.trainer.model
         data_parallel = registry.get("data_parallel") or registry.get("distributed")
         fp16_scaler = getattr(self.trainer, "scaler", None)
@@ -574,6 +575,15 @@ def save(self, update, iteration=None, update_best=False):
         with open_if_main(current_ckpt_filepath, "wb") as f:
             self.save_func(ckpt, f)
 
+        # Save the current checkpoint as W&B artifacts for model versioning.
+        if self.config.training.wandb.log_checkpoint:
+            logger.info(
+                "Saving current checkpoint as W&B Artifacts for model versioning"
+            )
+            self.trainer.logistics_callback.wandb_logger.log_model_checkpoint(
+                current_ckpt_filepath
+            )
+
         # Remove old checkpoints if max_to_keep is set
         # In XLA, only delete checkpoint files in main process
         if self.max_to_keep > 0 and is_main():

mmf/utils/logger.py

@@ -225,6 +225,12 @@ def summarize_report(
     if not is_main() and not is_xla():
         return
 
+    # Log the learning rate if available
+    if wandb_logger and "lr" in extra:
+        wandb_logger.log_metrics(
+            {"train/learning_rate": float(extra["lr"])}, commit=False
+        )
+
     if tb_writer:
         scalar_dict = meter.get_scalar_dict()
         tb_writer.add_scalars(scalar_dict, current_iteration)
@@ -395,21 +401,19 @@ class WandbLogger:
     Log using `Weights and Biases`.
 
     Args:
-        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`.
+        entity: An entity is a username or team name where you're sending runs.
+        config: Configuration for the run.
+        project: Name of the W&B project.
 
     Raises:
         ImportError: If wandb package is not installed.
     """
 
     def __init__(
         self,
-        name: Optional[str] = None,
-        save_dir: Optional[str] = None,
+        entity: Optional[str] = None,
+        config: Optional[Dict] = None,
         project: Optional[str] = None,
-        **init_kwargs,
     ):
         try:
             import wandb
@@ -420,10 +424,13 @@ def __init__(
             )
 
         self._wandb = wandb
-
-        self._wandb_init = dict(name=name, project=project, dir=save_dir)
-
-        self._wandb_init.update(**init_kwargs)
+        self._wandb_init = dict(entity=entity, config=config, project=project)
+        wandb_kwargs = dict(config.training.wandb)
+        wandb_kwargs.pop("enabled")
+        wandb_kwargs.pop("entity")
+        wandb_kwargs.pop("project")
+        wandb_kwargs.pop("log_checkpoint")
+        self._wandb_init.update(**wandb_kwargs)
 
         self.setup()
 
@@ -453,14 +460,33 @@ 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, commit=commit)
+
+    def log_model_checkpoint(self, model_path):
+        """
+        Log the model checkpoint to the wandb dashboard.
+
+        Args:
+            model_path (str): Path to the model file.
         """
         if not self._should_log_wandb():
             return
 
-        self._wandb.log(metrics)
+        model_artifact = self._wandb.Artifact(
+            "run_" + self._wandb.run.id + "_model", type="model"
+        )
+
+        model_artifact.add_file(model_path, name="current.pt")
+        self._wandb.log_artifact(model_artifact, aliases=["latest"])

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
+        enabled: true
+        # 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}
+        project: mmf
         # Experiment/ run name to be used while logging the experiment
         # under the project with wandb
-        wandb_runname: ${training.experiment_name}
+        name: ${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.
 
-`training.wandb.enabled=True`
+* 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. The default project name is `mmf` and the default run name is `${training.experiment_name}`.
+
+    `training.wandb.project=<ProjectName>` <br />
+    `training.wandb.name=<RunName>`
+
+* To change the path to the directory where wandb metadata would be stored (Default: `env.log_dir`):
+
+    `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. Make sure to use the same key name in the config file as defined in the documentation.
 
-`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