Learning Data Augmentation Strategies for Object Detection

TL;DR

This work introduces a reinforcement-learning-based search to automatically discover data augmentation policies tailored for object detection, incorporating bounding-box–aware transformations. The learned policies significantly improve mAP on COCO (up to +2.3 mAP) and even push a single model to state-of-the-art performance (50.7 mAP) with larger input sizes and anchors, while transferring effectively to other datasets like PASCAL VOC and across different backbones and detectors. The results show that color, geometric, and bbox-specific augmentations collectively drive the gains, especially on small datasets and small objects, and that the learned policy provides regularization without relying on other regularizers. Overall, the approach reduces the need for additional labeled data and demonstrates broad applicability across detection settings.

Abstract

Data augmentation is a critical component of training deep learning models. Although data augmentation has been shown to significantly improve image classification, its potential has not been thoroughly investigated for object detection. Given the additional cost for annotating images for object detection, data augmentation may be of even greater importance for this computer vision task. In this work, we study the impact of data augmentation on object detection. We first demonstrate that data augmentation operations borrowed from image classification may be helpful for training detection models, but the improvement is limited. Thus, we investigate how learned, specialized data augmentation policies improve generalization performance for detection models. Importantly, these augmentation policies only affect training and leave a trained model unchanged during evaluation. Experiments on the COCO dataset indicate that an optimized data augmentation policy improves detection accuracy by more than +2.3 mAP, and allow a single inference model to achieve a state-of-the-art accuracy of 50.7 mAP. Importantly, the best policy found on COCO may be transferred unchanged to other detection datasets and models to improve predictive accuracy. For example, the best augmentation policy identified with COCO improves a strong baseline on PASCAL-VOC by +2.7 mAP. Our results also reveal that a learned augmentation policy is superior to state-of-the-art architecture regularization methods for object detection, even when considering strong baselines. Code for training with the learned policy is available online at https://github.com/tensorflow/tpu/tree/master/models/official/detection

Figures (6)

• Figure 1: Learned augmentation policy systematically improves object detection performance. Left: Learned augmentation policy applied to example from COCO dataset lin2014microsoft. Right: Mean average precision for RetinaNet lin2017focal with a ResNet-50 backbone on COCO lin2014microsoft with and without learned augmentation policy (red and black, respectively).
• Figure 2: Examples of learned augmentation sub-policies. 5 examples of learned sub-policies applied to one example image. Each column corresponds to a different random sample of the corresponding sub-policy. Each step of an augmentation sub-policy consists of a triplet corresponding to the operation, the probability of application and a magnitude measure. The bounding box is adjusted to maintain consistency with the applied augmentation. Note the probability and magnitude are discretized values (see text for details).
• Figure 3: Percentage improvement in mAP for objects of different sizes due to the learned augmentation policy.
• Figure 4: Percentage improvement due to the learned augmentation policy on mAP, AP50, and AP75, relative to models trained with baseline augmentation.
• Figure 5: Training loss vs. number of training examples for baseline model (black) and with the learned augmentation policy (red).