Weights & Biases handler for MonAI

[soumik12345]

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

• 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).
• 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.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[Nic-Ma]

I see the wandb is global, is it thread-safe and supports multi-processing?
I would suggest to add some tests to call wandb.init(dir=tempdir.name) and the whole handler in multi-thread and multi-processing logic.
For multi-processing test, you can refer to:
https://github.com/Project-MONAI/MONAI/blob/dev/tests/test_cumulative_average_dist.py
For multi-thread test, you can easily add a test function here, refer to:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/dataset.py#L875-L878
What do you think?

Thanks.

[binliunls]

I cannot open the ignite concept link and the tutorial link below. May need to update these links.

Thanks,
Bin

[binliunls]

In the doc-string, it says that the epoch_event_writer: ..... Must accept the parameter "engine" and "summary_writer" ..... Whereas, I don't see the summary_writer parameter here. It's a little bit confusing for me. Could you please explain it?

Thanks,
Bin

[soumik12345]

Its a typo from an older draft of the PR, where I was mimicking the behavior of the Tensorboard Stats handler. I will update the docstring.

[binliunls]

Shoud a dict parameter like hyperparameters be added to the init function to record some users defined parameters with wandb.config?

Thanks,
Bin

[Nic-Ma]

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

Thanks.

[binliunls]

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()

[Nic-Ma]

Hi @soumik12345 ,

Do you still plan to complete this PR?

Thanks.

[soumik12345]

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 :)

[DanielNobbe]

Is this still being worked on?

[ericspod]

Hi @DanielNobbe sorry about this PR, it has seemingly fallen through the cracks quite a while ago. @soumik12345 I realised this is rather old now but would you be able to update the PR so we can merge this? Looking at it now I think we're fine to merge this if we can resolve the conflicts and run the tests. Thanks!