💡 Train an Object Detection+Semantic Segmentation Model with MMDetection and Weights and Biases

Weights and Biases

In this colab, we will train a Mask RCNN model using MMDetection on a small Balloon dataset.Through this colab you will learn to:

• use MMDetection to train an object detection + Semantic Segmentation model on a custom dataset,
• use Weights and Biases to log training and validation metrics, visualize model predictions, visualize raw validation dataset, and more.

This colab in particular, will showcase a dedicated MMDetWandbHook for MMDetection that can be used to:

✅ Log training and evaluation metrics.
✅ Log versioned model checkpoints.
✅ Log versioned validation dataset with ground truth bounding boxes and segmentation masks.
✅ Log and visualize model predictions.

But before we continue, here's a quick summary of MMDetection and W&B if you are not familiar with them.

📸 MMDetection

MMDetection is an open source object detection toolbox based on PyTorch. It provides composable components that are easy to customize and has out-of-box support for single and multi GPU training/inference. It also has hundreds of pretrained detection models in Model Zoo, and supports multiple standard datasets. Check out the GitHub repository here.

📸 Weights and Biases

Consider Weights and Biases (W&B) to be the GitHub for machine learning. Use W&B for machine learning experiment tracking, dataset and model versioning, project collaboration, hyperparameter optimization, dataset exploration, model evaluation and so much more. If you are new to W&B, check out this intro colab.

⚽️ Imports and Setup

1️⃣ Install MMDetection

MMDetection is heavily dependent on the MMCV library. We will have to install the version of MMCV that is compatible with the given PyTorch version. Check out the Installation documentation for more details.

2️⃣ Install Weights and Biases

Install the latest version of W&B.

3️⃣ General Imports

4️⃣ Login with your W&B account

Create a free W&B account (it's free for personal and academic usage). Create a new API key at wandb.ai/settings and store it securely. API keys can only be viewed once when created.

🏀 Dataset

We will be using a small Balloon dataset for this colab notebook.

1️⃣ Download the dataset

Note: The train and val folders contain the images along with annotations as .json files.

2️⃣ Convert annotations to COCO format

To support a new data format, it's recommended to convert the annotations to COCO format or PASCAL VOC format.

If you are converting annotations to COCO format, do so offline and use the CocoDataset class. If you are converting it to the PASCAL format, use the VOCDataset class. You will see the usage below.

You can find more details about customizing the dataset here.

🏈 Model

There are over hundred pre-trained object detectors provided by MMDetection via Model Zoo. Check out the Model Zoo documentation page.

You can also customize the model's backbone, neck, head, ROI, and loss. More on customizing the model here.

1️⃣ Download the model

We will be using a pretrained model checkpoint to fine tune on our custom dataset. Let's download the model in the checkpoints directory.

I am using the this model with this config from the Model Zoo. You can find similar Mask RCNN model here.

⚾️ Configuration

MMDetection relies heavily on a config system. In the cell below, we will be loading a config file and modify few of the methods as per the need of this notebook.

Note that both train and test dataloaders will use the same training samples. This is not a recommended practice but for the sake of a simplified notebook, let's use it.

Learn more about the MMDetection Config system here.

1️⃣ Load the config file

2️⃣ Modify data config

3️⃣ Modify model config

4️⃣ Modify training config

5️⃣ Modify evaluation config

Note that we are evaluating the model on the validation dataset after every 2 epochs while saving the checkpoints after every 4 epochs.

⭐️ If you want to log the configuration to W&B for efficient experiment tracking, add the config filename to a dict with key exp_name.

🎾 Define Weights and Biases Hook

MMDetection comes with a dedicated Weights and Biases Hook - MMDetWandHook.

With this dedicated hook, you can:

• log train and eval metrics along with system (CPU/GPU) metrics,
• visualize the validation dataset as interactive W&B Tables,
• visualize the model prediction as interactive W&B Tables, and
• save the model checkpoints as W&B Artifacts.

To use this hook, you can append a dict to log_config.hooks. The log_config wraps multiple logger hooks like the TextLoggerHook used below.

There are four important arguments in the MMDetWandbHook that can help you get the most out of MMDetection.

• init_kwargs: Use this argument to in-turn pass arguments to wandb.init. You can use it to set the W&B project name, set the team name (entity) if you want to log the runs to a team account, pass the configuration, and more. Check out the arguments that you can pass to wandb.init here.

• log_checkpoint: Save the checkpoint at every checkpoint interval as W&B Artifacts. Use this for model versioning where each version is a checkpoint. The checkpoint interval is set using checkpoint_config.interval (starred above). Note that this feature is dependent on MMCV's CheckpointHook.

• log_checkpoint_metadata: Log the evaluation metrics computed on the validation data with the checkpoint, along with current epoch as a metadata to that checkpoint.

• num_eval_images: At every evaluation interval, the MMDetWandbHook logs the model prediction as interactive W&B Tables. The eval interval is determined by evaluation.interval (starred above). The number of samples logged is given by num_eval_images. Before training begins, the callback logs the validation data (images along with bounding box and masks). At every evaluation interval, the model predictions are logged. This Feature is dependent on MMCV's EvalHook or DistEvalHook.

• bbox_score_thr: Threshold for bounding box scores.

🏐 Train

Now that we have the dataset, pretrained model weight, and have defined the configs. Let's stitch them together to train an object detector.

1️⃣ Build the Dataset

2️⃣ Build the Model

3️⃣ Train with W&B

4️⃣ Notes on using MMDetWandbHook.

Using MMDetWandbHook is easy and in most cases it will throw friendly UserWarning if something is not quite right. However in the best interest, here are some of things and best practices you should keep in mind:

• The MMDetWandbHook depends on CheckpointHook for logging the checkpoints as W&B Artifacts and EvalHook/DistEvalHook for logging validation data and model predictions. If anyone or both aren't available, this hook will give UserWarning and not cause any error.

• The priority of both CheckpointHook and EvalHook/DistEvalHook should be more than MMDetWandbHook.

• The validation data is logged once as val_data W&B Table. The evaluation tables, use reference to this data thus you will not be uploading the same data multiple times.

• If you want to log the configuration to W&B, pass this key-value pair 'config': cfg._cfg_dict.to_dict() to init_kwargs. However, it's recommended to use meta['exp_name'] = config_filename and pass meta to train_detector.