Weight & Biases support

🚀 Feature

Adding support for wandb library

Motivation

W&B is one of the most advanced and feature-rich experiment tracking tools. It would be great to have support in detectron2. Great example: https://app.wandb.ai/syllogismos/keras-retinanet

Pitch

I implemented it in my own project as wanbb_writer.py module

import wandb
from detectron2.utils.events import (
    EventWriter,
    get_event_storage,
)


class WAndBWriter(EventWriter):
    """
    Write all scalars to a wandb tool.
    """

    def __init__(self, window_size: int = 20):
        self._window_size = window_size

    def write(self):
        storage = get_event_storage()
        for k, v in storage.latest_with_smoothing_hint(self._window_size).items():
            wandb.log({f"{k}": v}, step=storage.iter)

        if len(storage.vis_data) >= 1:
            for img_name, img, step_num in storage.vis_data:
                self._writer.add_image(img_name, img, step_num)
            storage.clear_images()

    def close(self):
        pass

and add this class as one of the writer here https://github.com/facebookresearch/detectron2/blob/master/tools/plain_train_net.py#L136

If someone finds this feature useful as I do, I would be happy to send a PR with W&B writer and tests for it.

cc: @cvphelps

[ppwwyyxx]

I suppose wandb will be enabled by just adding import wandb; wandb.init(sync_tensorboard=True) to your code. It doesn’t look like we need to do anything for it.

AP curves will be available if you set TEST.EVAL_PERIOD.

[marcosterland]

Probably something has changes meanwhile, but the example of @truskovskiyk doesn’t work for me. Below is my adaptation, which works with v0.3. I’ve noticed W&B has issues with / in the metric names.

class WAndBWriter(EventWriter):
    """
    Write all scalars to a wandb tool.
    """

    def __init__(self, window_size: int = 20):
        self._window_size = window_size

    def write(self):
        storage = get_event_storage()
        stats = {}
        for k, v in storage.latest_with_smoothing_hint(self._window_size).items():
            stats[k.replace("/", "-")] = v[0]
        wandb.log(stats, step=storage.iter)

        if len(storage._vis_data) >= 1:
            for img_name, img, step_num in storage._vis_data:
                self._writer.add_image(img_name, img, step_num)
            storage.clear_images()

    def close(self):
        pass