Skip to content
Snippets Groups Projects
Commit 46ff0937 authored by Manon Blanco's avatar Manon Blanco Committed by Solene Tarride
Browse files

Weights and Biases integration during training - Initialisation

parent 646805ae
No related branches found
No related tags found
1 merge request!440Weights and Biases integration during training - Initialisation
......@@ -8,6 +8,7 @@
"aws_access_key_id": "",
"aws_secret_access_key": ""
},
"wandb": {},
"dataset": {
"datasets": {
"$dataset_name": "$dataset_path"
......
......@@ -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"
......
......@@ -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` | |
......@@ -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"]
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment