ResNeSt: Split-Attention Networks

[Paper] [GitHub] [PyTorch] [Gluon]

Motivation

While image classification models have kept advancing, most downstream applications such as object detection and semantic segmentation still rely on the ResNet variants as the backbone, because of its simple and modular structure. However, the vanilla ResNet is not designed for downstream tasks, since the receptive-field size and the lack of cross-channel interaction are known to be limiting. Boosting the backbone accuracy usually transfers directly to downstream gains, so we set out to design a backbone that keeps the ResNet layout but improves feature representation.

Recent works such as SE-Net and SK-Net have shown that channel-wise attention and feature-map attention can improve the representation power. Meanwhile, ResNeXt introduced multi-path feature representations via grouped convolutions. We ask: can we combine these ideas — multi-path representation and channel-wise attention — inside a single, ResNet-friendly building block?

Split-Attention Block

The core contribution of ResNeSt is the Split-Attention block, a computational unit that enables attention across feature-map groups.

Similar to ResNeXt, the feature can be divided into several groups, and the number of feature-map groups is given by a cardinality hyperparameter \(K\). We refer to the resulting feature-map groups as cardinal groups. We introduce a new radix hyperparameter \(R\) that indicates the number of splits within a cardinal group, so the total number of feature groups is \(G = KR\).

For each cardinal group, a combined representation is obtained by fusing the split representations via an element-wise summation. The global contextual information is gathered with global average pooling along the spatial dimension. A soft assignment weight is then computed by applying two fully-connected layers with a softmax (across splits within the same cardinal group) — this is the “split attention”. The output of each cardinal group is a weighted fusion of the splits using these attention weights. The outputs from all cardinal groups are concatenated along the channel dimension, and a shortcut connection is added to form the block output, matching the ResNet residual style.

Split-Attention generalizes prior designs: it reduces to SE-Net attention when \(R=1\), and to SK-Net when \(K=1, R=2\) (with slight differences in the softmax formulation).

Network Design

The overall ResNeSt architecture follows the ResNet stage layout, and simply replaces each residual bottleneck block with a Split-Attention block. This gives us drop-in replacements for ResNet-50 / 101 / 200 / 269, which we call ResNeSt-50 / 101 / 200 / 269. Because the block preserves ResNet’s input/output signature, ResNeSt can be used as a backbone for downstream detection and segmentation frameworks without any pipeline changes.

We also adopt a few training tricks that have been shown effective for training image classification models, including large mini-batch distributed training, cosine learning-rate schedule, label smoothing, auto augmentation, mixup training, large crop size and regularization.

ImageNet Results

ResNeSt outperforms other networks with similar model complexities. ResNeSt-50 achieves 81.13% top-1 accuracy on ImageNet with the crop size of 224x224, outperforming previous best ResNet variant by more than 1%. This improvement may benefit many downstream applications.

Transfer Learning Results

Simply plugging ResNeSt into existing frameworks brings large gains on downstream tasks.

• Object detection. With Faster-RCNN and Cascade-RCNN on MS-COCO, swapping the ResNet-50/101 backbone with ResNeSt-50/101 improves mAP by around 3% (e.g. Faster-RCNN box mAP goes from 39.3 → 42.3 with ResNeSt-50).
• Instance segmentation. With Mask-RCNN and Cascade-Mask-RCNN, ResNeSt improves both box and mask mAP by a similar margin.
• Semantic segmentation. On ADE20K, DeepLabV3 with a ResNeSt-101 backbone achieves 46.9% mIoU (single-scale), outperforming DeepLabV3 with a ResNet-101 backbone by more than 1%.

Usage

The pretrained models can be loaded directly from PyTorch Hub in a couple of lines:

import torch
# get list of models
torch.hub.list('zhanghang1989/ResNeSt', force_reload=True)
# load pretrained models, using ResNeSt-50 as an example
net = torch.hub.load('zhanghang1989/ResNeSt', 'resnest50', pretrained=True)

Or install via pip and use directly:

pip install resnest --pre

from resnest.torch import resnest50
net = resnest50(pretrained=True)

Please refer to the GitHub repo for training scripts and reproducing the results.

Citation

@article{zhang2020resnest,
    title={ResNeSt: Split-Attention Networks},
    author={Zhang, Hang and Wu, Chongruo and Zhang, Zhongyue and Zhu, Yi and Lin, Haibin
        and Zhang, Zhi and Sun, Yue and He, Tong and Mueller, Jonas and Manmatha, R.
        and Li, Mu and Smola, Alexander},
    journal={arXiv preprint arXiv:2004.08955},
    year={2020}
}