Weights and Biases integration during training - Initialisation

configs/quickstart.json

@@ -8,6 +8,7 @@
         "aws_access_key_id": "",
         "aws_secret_access_key": ""
     },
+    "wandb": {},
     "dataset": {
         "datasets": {
             "$dataset_name": "$dataset_path"

dan/ocr/train.py

@@ -10,6 +10,7 @@ from copy import deepcopy
 import numpy as np
 import torch
 import torch.multiprocessing as mp
+import wandb
 
 from dan.ocr.manager.training import Manager
 from dan.ocr.mlflow import MLFLOW_AVAILABLE
@@ -106,6 +107,14 @@ def run(config: dict):
     """
     Main program, training a new model, using a valid configuration
     """
+    # Start "Weights & Biases" as soon as possible
+    if config.get("wandb"):
+        wandb_config = config["wandb"].pop("config", {})
+        wandb.init(
+            **config["wandb"],
+            config={**config, **wandb_config},
+        )
+
     names = list(config["dataset"]["datasets"].keys())
     # We should only have one dataset
     assert len(names) == 1, f"Found {len(names)} datasets but only one is expected"

docs/usage/train/config.md

@@ -271,3 +271,19 @@ $ pip install .[mlflow]
 | `mlflow.experiment_id`         | ID of the current experiment in MLFlow. | `str` |         |
 | `mlflow.aws_access_key_id`     | Access key ID to the AWS server.        | `str` |         |
 | `mlflow.aws_secret_access_key` | Secret access key to the AWS server.    | `str` |         |
+
+## Weights & Biases logging
+
+To log your experiment on [Weights & Biases](https://wandb.ai/), you need to:
+
+- [login to Weights & Biases](https://docs.wandb.ai/ref/cli/wandb-login) via
+
+```shell
+wandb login
+```
+
+- update the following arguments:
+
+| Name    | Description                                                                                                                                                                          | Type   | Default |
+| ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------ | ------- |
+| `wandb` | Key and values to use to initialise your experiment on Weights & Biases. See the full list of available keys on [the official documentation](https://docs.wandb.ai/ref/python/init). | `dict` |         |

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "torchaudio==2.2.0",
     "torchvision==0.17.0",
     "tqdm==4.65.0",
+    "wandb==0.17.0",
 ]
 
 keywords = ["python", "HTR", "OCR", "NER", "machine learning", "pytorch"]