MONAI Weights & Biases handler for MonAI

Features Contributed

WandbStatsHandler defines a set of Ignite Event-handlers for all the Weights & Biases logging logic. It can be used for any Ignite Engine(trainer, validator, and evaluator) and support both epoch level and iteration level. The expected data source is Ignite engine.state.output and engine.state.metrics. Default behaviors:
- When EPOCH_COMPLETED, write each dictionary item in engine.state.metrics to Weights & Biases.
- When ITERATION_COMPLETED, write each dictionary item in self.output_transform(engine.state.output) to Weights & Biases.

The following colab notebook and Weights & Biases run demonstrate the usage of these handlers and their results respectively:

Some additional Weights & Biases features:

When TensorBoardStatsHandler and TensorBoardImageHandler are used inside a wandb run, Weights & Biases automatically hosts the Tensorboard instance inside the run if during wandb.init(), sync_tensorboard is set to True.
When used with TensorBoardImageHandler, the images and videos are automatically logged to Weights & Biases media panel if during wandb.init(), sync_tensorboard is set to True.

Types of changes

[x] Non-breaking change (fix or new feature that would not break existing functionality).
[ ] Breaking change (fix or new feature that would cause existing functionality to change).
[x] New tests added to cover the changes.
[ ] Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
[ ] Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
[x] In-line docstrings updated.
[x] Documentation updated, tested make html command in the docs/ folder.

May 16 '23 07:05 soumik12345

@SachidanandAlle Can help share more concerns about multi-thread running of the hander.

Thanks.

May 19 '23 03:05 Nic-Ma

When running with the multithread code shown below, the wandb handler records the thread results into one run as shown in the pic. However, I think several indiviual runs should be created to record different threads.

from __future__ import annotations

import unittest
import wandb
import torch
from concurrent.futures import ThreadPoolExecutor

from ignite.engine import Engine

from monai.handlers import WandbStatsHandler


def dummy_train(start):
    # set up engine
    def _train_func(engine, batch):
        return batch + 1.0

    engine = Engine(_train_func)

    # set up testing handler
    handler = WandbStatsHandler(
        output_transform=lambda x: x,
    )
    handler.attach(engine)
    engine.run(torch.tensor([start]), max_epochs=5)


class TestHandlerWB(unittest.TestCase):
    def test_multi_thread(self):
        wandb.init(
            project="multithread-handlers", save_code=True, sync_tensorboard=True
        )
        with ThreadPoolExecutor(2, "Training") as executor:
            for t in range(2):
                executor.submit(dummy_train, t + 2)


if __name__ == "__main__":
    unittest.main()

May 22 '23 07:05 binliunls

Hi @soumik12345 ,

Do you still plan to complete this PR?

Thanks.

Aug 08 '23 14:08 Nic-Ma

Hi @soumik12345 ,

Do you still plan to complete this PR?

Thanks.

Hi @Nic-Ma I will be completing this PR soon, thanks for being patient :)

Aug 08 '23 15:08 soumik12345