diff --git a/configs/quickstart.json b/configs/quickstart.json index 648fa5f6071cfd90004ae1b8339080f2c63614bb..a9fa734e2bb38c2686bb779800520e33d7ad9c63 100644 --- a/configs/quickstart.json +++ b/configs/quickstart.json @@ -8,6 +8,7 @@ "aws_access_key_id": "", "aws_secret_access_key": "" }, + "wandb": {}, "dataset": { "datasets": { "$dataset_name": "$dataset_path" diff --git a/dan/ocr/train.py b/dan/ocr/train.py index c691c0e267c7a2726c61da9e251f2043e35a76e9..6f64ef5bce9a3254d8f28ae4f5a5e0b29a2dda68 100644 --- a/dan/ocr/train.py +++ b/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" diff --git a/docs/usage/train/config.md b/docs/usage/train/config.md index 14a145fbfd3a209498688a334cb27d4fef2bcad0..fc6765ddea282d4db30f248a11e14cb669f687d9 100644 --- a/docs/usage/train/config.md +++ b/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` | | diff --git a/pyproject.toml b/pyproject.toml index 723eb2b0a868cfcc22b952ec668f744524a79b01..5e636fafb84a546904197fb74ae1ccec08a10d40 100644 --- a/pyproject.toml +++ b/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"]