Train FasterRCNN faster with 16-bit precision in Detectron2

One of the features of Detectron2 is that it is faster than its previous versions. Now, with the release of the wonderful library PyTorch Lightning, it is possible to train models using float16, bfloat16, float32, and other formats. For running large models, accommodating a large batch size, or faster training, it is advisable to use Automatic Mixed Precision in PyTorch.

Using PyTorch Lightning

This can be easily done in PyTorch Lightning by setting the precision argument to ‘16’ or ‘bf16’ when instantiating the Trainer. Detectron2 provides the provision for training FasterRCNN using PyTorch Lightning by providing a training script lightning_train_net.py in the detectron2/tree/main/tools directory.

At the time of writing this article (2nd Sept 2023), there are some minor bugs in the lightning_train_net.py script that will throw an error when running without rectifying them. To avoid those errors, please make the changes as reported in this Issues thread.

Without using PyTorch Lightning

Detectron2 also provides the provision for automatic mixed precision without PyTorch Lighting. Detectron2 provides an AMPTrainer class in detectron2/blob/main/detectron2/engine/train_loop.py. Thus, using the training script train_net.py, we can train Object Detection or Segmentation models with 16-bit precision.

How to do it?

To train a FasterRCNN, MaskRCNN, or any other models available in Detectron2 MODEL_ZOO with automatic mixed precision, we need to set the argument SOLVER.AMP.ENABLED to True.

With PyTorch Lightning

This sets the precision argument in PyTorch Lightning Trainer to ‘16’ for 16-bit precision in lightning_train_net.py. This value can be changed to ‘bf16’ in lightning_train_net.py for bf16 precision training if supported.

Without PyTorch Lightning

In train_net.py, however, setting the above argument to True sets the trainer to an AMPTrainer instance, and a SimpleTrainer instance if set to False.

Other ways to do the same thing

Otherwise, this argument can also be set in the ‘.yaml’ config file by directly adding the following lines under SOLVER

SOLVER
  .
  .
  .  
  AMP
    ENABLED: True

Finally

The whole configuration file given here for PASCAL VOC Object Detection will look like this when modified

_BASE_: "../Base-RCNN-C4.yaml"
MODEL:
  WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl"
  MASK_ON: False
  RESNETS:
    DEPTH: 50
  ROI_HEADS:
    NUM_CLASSES: 20
INPUT:
  MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
  MIN_SIZE_TEST: 800
DATASETS:
  TRAIN: ('voc_2007_trainval', 'voc_2012_trainval')
  TEST: ('voc_2007_test',)
SOLVER
  STEPS: (12000, 16000)
  MAX_ITER: 18000  # 17.4 epochs
  WARMUP_ITERS: 100
  IMS_PER_BATCH: 16
  BASE_LR: 0.1
  AMP
    ENABLED: True

Clap and share if you like this article. Follow for more.