MMSegmentation 中的 ResNeSt 分割骨干网络:Split-Attention 原理、源码解析与 Cityscapes/ADE20K 实战配置
【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation
本文基于 MMSegmentation 仓库 configs/resnest/README.md 与 mmseg/models/backbones/resnest.py 编写。ResNeSt(Split-Attention Networks)是一种将通道注意力(channel-wise attention)与多路径表示(multi-path representation)模块化的骨干网络,MMSegmentation 将其作为 FCN、PSPNet、DeepLabV3、DeepLabV3+ 等多种分割头的高性能 backbone 接入,并提供了 Cityscapes 与 ADE20K 上的完整训练配置与预训练权重。读完本文,你将掌握 ResNeSt 的 Split-Attention 计算块原理、MMSegmentation 中 ResNeSt 的源码实现细节,以及如何基于仓库配置一键复现训练与推理。
一、ResNeSt 是什么:论文背景与核心思想
ResNeSt(Split-Attention Networks)出自论文ResNeSt: Split-Attention Networks(arXiv:2004.08955,Zhang 等人,2020)。其核心观察是:特征图注意力(featuremap attention)与多路径表示(multi-path representation)对视觉识别任务至关重要。ResNeSt 将这两者结合,提出一个模块化的计算单元——Split-Attention 块:
- 在不同的网络分支(branches)上施加通道维度的注意力(channel-wise attention),以捕捉跨特征交互(cross-feature interactions);
- 让不同分支学习多样化的表示(diverse representations);
- 整个设计最终收敛为一个简单、统一的计算块,仅需少量变量即可参数化。
原论文报告:ResNeSt 在图像分类上以更优的精度/延迟权衡超越 EfficientNet,作为骨干网络在多个公开基准上取得了优秀的迁移学习结果,并被 COCO-LVIS 挑战赛的获胜方案所采用。
论文信息与引用
该论文对应的 BibTeX 引用(来自 configs/resnest/README.md 的 Citation 小节):
@article{zhang2020resnest, title={ResNeSt: Split-Attention Networks}, author={Zhang, Hang and Wu, Chongruo and Zhang, Zhongyue and Zhu, Yi and Zhang, Zhi and Lin, Haibin and Sun, Yue and He, Tong and Muller, Jonas and Manmatha, R. and Li, Mu and Smola, Alexander}, journal={arXiv preprint arXiv:2004.08955}, year={2020} }二、MMSegmentation 中的 ResNeSt:源码级原理剖析
ResNeSt 在 MMSegmentation 中的实现位于 mmseg/models/backbones/resnest.py,由三个核心组件构成:RSoftmax、SplitAttentionConv2d与Bottleneck,并通过@MODELS.register_module()注册为ResNeSt,继承自 ResNet 的 V1d 变体(见 mmseg/models/backbones/resnet.py 中ResNetV1d)。
2.1 RSoftmax:Radix Softmax 模块
RSoftmax(resnest.py 第 16-37 行)是 Split-Attention 中产生注意力权重的激活函数,行为取决于radix参数:
- 当
radix > 1时,将输入重塑为(batch, groups, radix, -1)并转置为(batch, radix, groups, -1),在 radix 维度上做softmax,实现多个分支间的竞争归一化; - 当
radix == 1时,退化为sigmoid激活(等价于 SE-Net 式的通道门控)。
class RSoftmax(nn.Module): def forward(self, x): batch = x.size(0) if self.radix > 1: x = x.view(batch, self.groups, self.radix, -1).transpose(1, 2) x = F.softmax(x, dim=1) x = x.reshape(batch, -1) else: x = torch.sigmoid(x) return x2.2 SplitAttentionConv2d:Split-Attention 卷积单元
SplitAttentionConv2d(resnest.py 第 40-144 行)是 ResNeSt 的核心计算块,其构造参数包括:
| 参数 | 类型 | 默认值 | 说明 |
|---|---|---|---|
in_channels | int | 必填 | 输入通道数,与nn.Conv2d一致 |
channels | int | 必填 | 输出通道数 |
kernel_size | int/tuple | 必填 | 卷积核大小 |
stride | int/tuple | 1 | 步长 |
padding | int/tuple | 0 | 填充 |
dilation | int/tuple | 1 | 膨胀率 |
groups | int | 1 | 分组数 |
radix | int | 2 | SplitAtConv2d 的分支基数 |
reduction_factor | int | 4 | inter_channels的缩减因子 |
conv_cfg | dict | None | 卷积层配置,默认使用普通 conv2d |
norm_cfg | dict | dict(type='BN') | 归一化层配置 |
dcn | dict | None | DCN(可变形卷积)配置 |
其前向过程(resnest.py 第 118-144 行)完整实现了 Split-Attention 的标准四步流程:
- Split(分支卷积):先经过
groups * radix分组的卷积self.conv,输出channels * radix个通道;当radix > 1时,将输出按 radix 拆分为多个分支(splits),并对所有分支求和得到gap; - SE 式压缩:对
gap做全局自适应平均池化(F.adaptive_avg_pool2d(gap, 1)); - 通道注意力生成:依次经过
fc1(将通道数压缩到inter_channels)、BN + ReLU、fc2(恢复为channels * radix),再经RSoftmax得到注意力权重; - 加权融合(Split Attention):当
radix > 1时,将注意力权重按 radix 拆分后与各分支逐元素相乘并求和(torch.sum(attens * splits, dim=1));当radix == 1时,直接与原始特征相乘。
其中inter_channels的计算式为max(in_channels * radix // reduction_factor, 32),即压缩后通道数至少为 32,保证小通道数场景下注意力瓶颈不会过窄。此外,该模块支持 DCN(可变形卷积):当传入dcn且未设置fallback_on_stride时,会断言conv_cfg必须为 None,并将dcn作为卷积配置(resnest.py 第 79-84 行)。
2.3 Bottleneck:集成 Split-Attention 的残差块
ResNeSt 的Bottleneck(resnest.py 第 147-267 行)继承自 ResNet 的Bottleneck(expansion = 4),关键差异在于:
- conv2 被替换为
SplitAttentionConv2d,并传入groups、radix、reduction_factor等参数(resnest.py 第 201-213 行); - avg_down_stride:当启用且
conv2_stride > 1时,stride 不再放在 3x3 卷积中,而是在卷积之后插入一个nn.AvgPool2d(3, conv2_stride, padding=1)下采样(resnest.py 第 185、216-217、241-242 行); - 支持
with_cp(checkpoint)以节省显存:当x.requires_grad时使用torch.utils.checkpoint包装内层前向(resnest.py 第 260-263 行)。
class Bottleneck(_Bottleneck): expansion = 4 def __init__(self, inplanes, planes, groups=1, base_width=4, base_channels=64, radix=2, reduction_factor=4, avg_down_stride=True, **kwargs): super().__init__(inplanes, planes, **kwargs) # ... self.conv2 = SplitAttentionConv2d( width, width, kernel_size=3, stride=1 if self.avg_down_stride else self.conv2_stride, padding=self.dilation, dilation=self.dilation, groups=groups, radix=radix, reduction_factor=reduction_factor, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, dcn=self.dcn)2.4 ResNeSt 骨干:注册、深度配置与 V1d 继承
ResNeSt类(resnest.py 第 270-318 行)通过@MODELS.register_module()注册到 MMSegmentation 的模型注册表中,可直接在配置中以type='ResNeSt'引用。
支持的深度(arch_settings,resnest.py 第 288-293 行):
| 深度 | 各阶段 Bottleneck 数量 |
|---|---|
| 50 | (3, 4, 6, 3) |
| 101 | (3, 4, 23, 3) |
| 152 | (3, 8, 36, 3) |
| 200 | (3, 24, 36, 3) |
构造参数:
| 参数 | 默认值 | 说明 |
|---|---|---|
groups | 1 | Bottleneck 中 3x3 卷积的分组数 |
base_width | 4 | 每组宽度(64x4d 表示groups=64, width_per_group=4) |
radix | 2 | SplitAttentionConv2d 的分支基数 |
reduction_factor | 4 | SplitAttentionConv2d 中间通道的缩减因子 |
avg_down_stride | True | 是否用平均池化实现下采样 |
| 其余 kwargs | - | 继承自 ResNet/ResNetV1d 的参数(depth、in_channels、stem_channels、out_indices、frozen_stages等) |
与 ResNetV1d 的关系:ResNeSt继承自ResNetV1d(resnet.py 第 703-712 行),因此天然具备 V1d 的两个特性:
- deep_stem:用三个 3x3 卷积(通道为
stem_channels//2 → stem_channels//2 → stem_channels)替换标准 ResNet 的单个 7x7 卷积(resnet.py 第 591-624 行); - avg_down:下采样残差块中先做 2x2 stride=2 的平均池化,再使用 stride=1 的卷积。
这解释了为何所有 ResNeSt 配置都将stem_channels设为128:ResNeSt 原始设计采用 64→128 的 stem 结构(相比 ResNet 的 64 通道 stem 更宽),从而与官方预训练权重的结构保持一致。
三、仓库配置解析:如何在 MMSegmentation 中启用 ResNeSt
3.1 配置复用模式
configs/resnest/目录下共 8 个训练配置,全部采用最小覆盖(base)继承模式:只替换 backbone 为ResNeSt,其余(分割头、数据集、调度、runtime)完全继承自对应算法的标准 ResNet-101 配置。例如:
- resnest_s101-d8_fcn_4xb2-80k_cityscapes-512x1024.py → 继承 configs/fcn/fcn_r101-d8_4xb2-80k_cityscapes-512x1024.py
- resnest_s101-d8_pspnet_4xb2-80k_cityscapes512x1024.py → 继承 configs/pspnet/pspnet_r101-d8_4xb2-80k_cityscapes-512x1024.py
- resnest_s101-d8_deeplabv3_4xb2-80k_cityscapes-512x1024.py → 继承 configs/deeplabv3/deeplabv3_r101-d8_4xb2-80k_cityscapes-512x1024.py
- resnest_s101-d8_deeplabv3plus_4xb2-80k_cityscapes-512x1024.py → 继承 configs/deeplabv3plus/deeplabv3plus_r101-d8_4xb2-80k_cityscapes-512x1024.py
- ADE20K 下的 4 个配置同理继承各自的
_4xb4-160k_ade20k-512x512版本
以 FCN + Cityscapes 为例,完整配置内容为:
_base_ = '../fcn/fcn_r101-d8_4xb2-80k_cityscapes-512x1024.py' model = dict( pretrained='open-mmlab://resnest101', backbone=dict( type='ResNeSt', stem_channels=128, radix=2, reduction_factor=4, avg_down_stride=True))3.2 关键配置项逐条说明
| 配置项 | 值 | 说明 |
|---|---|---|
model.pretrained | 'open-mmlab://resnest101' | 加载 OpenMMLab 托管的 ResNeSt-101 ImageNet 预训练权重 |
backbone.type | 'ResNeSt' | 注册表中的骨干类型(见 resnest.py 第 270 行) |
backbone.stem_channels | 128 | 输入 stem 通道数,与预训练权重结构匹配 |
backbone.radix | 2 | Split-Attention 分支基数 |
backbone.reduction_factor | 4 | 注意力瓶颈缩减因子 |
backbone.avg_down_stride | True | 用平均池化实现 stride 下采样 |
由于ResNeSt继承自ResNetV1d,depth、dilations(d8 即输出 stride 8)、out_indices等参数均沿用基类默认值,配置中的d8表示采用 dilation 策略使输出步长为 8,保证高分辨率分割特征。
四、实验结果:Cityscapes 与 ADE20K 基准
以下结果来自 configs/resnest/README.md 的 Results and models 小节,训练资源均为 4 块 V100 GPU。mIoU(ms+flip)表示多尺度 + 水平翻转测试的指标。
4.1 Cityscapes(512x1024,80k 迭代)
| Method | Backbone | Crop Size | Lr schd | Mem (GB) | Inf time (fps) | Device | mIoU | mIoU(ms+flip) |
|---|---|---|---|---|---|---|---|---|
| FCN | S-101-D8 | 512x1024 | 80000 | 11.4 | 2.39 | V100 | 77.56 | 78.98 |
| PSPNet | S-101-D8 | 512x1024 | 80000 | 11.8 | 2.52 | V100 | 78.57 | 79.19 |
| DeepLabV3 | S-101-D8 | 512x1024 | 80000 | 11.9 | 1.88 | V100 | 79.67 | 80.51 |
| DeepLabV3+ | S-101-D8 | 512x1024 | 80000 | 13.2 | 2.36 | V100 | 79.62 | 80.27 |
4.2 ADE20K(512x512,160k 迭代)
| Method | Backbone | Crop Size | Lr schd | Mem (GB) | Inf time (fps) | Device | mIoU | mIoU(ms+flip) |
|---|---|---|---|---|---|---|---|---|
| FCN | S-101-D8 | 512x512 | 160000 | 14.2 | 12.86 | V100 | 45.62 | 46.16 |
| PSPNet | S-101-D8 | 512x512 | 160000 | 14.2 | 13.02 | V100 | 45.44 | 46.28 |
| DeepLabV3 | S-101-D8 | 512x512 | 160000 | 14.6 | 9.28 | V100 | 45.71 | 46.59 |
| DeepLabV3+ | S-101-D8 | 512x512 | 160000 | 16.2 | 11.96 | V100 | 46.47 | 47.27 |
数据解读:在 Cityscapes 上,DeepLabV3(79.67 mIoU)与 DeepLabV3+(79.62 mIoU)明显优于 FCN(77.56 mIoU)与 PSPNet(78.57 mIoU);在 ADE20K 上,DeepLabV3+ 以 46.47 mIoU 领先,且四种方法在 ms+flip 测试下均有约 0.5~1 个点的提升。S-101-D8 指 Split-Attention ResNeSt-101、输出 stride 8 的骨干配置。
上述 8 个模型的权重与训练日志地址、批次规模(Cityscapes 为 4x2=8,ADE20K 为 4x4=16)等元数据均登记在 configs/resnest/metafile.yaml 中,可配合 MMSegmentation 的模型索引机制使用。
五、实战:训练、测试与推理
5.1 单机多卡训练
使用 tools/train.py 与仓库提供的 tools/dist_train.sh 启动训练:
bash tools/dist_train.sh configs/resnest/resnest_s101-d8_deeplabv3_4xb2-80k_cityscapes-512x1024.py 8其中第二个参数为 GPU 数量。训练配置的pretrained='open-mmlab://resnest101'会自动下载 ResNeSt-101 预训练权重用于 backbone 初始化。
5.2 测试与指标复现
使用 tools/test.py 测试,--out保存预测结果,--eval mIoU计算 mIoU 指标:
bash tools/dist_test.sh configs/resnest/resnest_s101-d8_deeplabv3_4xb2-80k_cityscapes-512x1024.py \ work_dirs/resnest_s101-d8_deeplabv3_4xb2-80k_cityscapes-512x1024/latest.pth \ 8 --out results.pkl --eval mIoU如需复现表中的mIoU(ms+flip),可结合 MMSegmentation 的多尺度 + 翻转测试评估流程(tools/test.py支持的--aug-test选项)。
5.3 单图推理
使用仓库自带的推理脚本 demo/image_demo.py 直接对单张图片进行推理:
python demo/image_demo.py demo/demo.png \ configs/resnest/resnest_s101-d8_deeplabv3_4xb2-80k_cityscapes-512x1024.py \ /path/to/checkpoint.pth六、单元测试验证:ResNeSt 的正确性保障
仓库在 tests/test_models/test_backbones/test_resnest.py 中提供了专门的单元测试,覆盖两大场景:
1. Bottleneck 结构与前向(test_resnest_bottleneck):
- 非法
style参数(如'tensorflow')会触发AssertionError; BottleneckS(64, 256, radix=2, reduction_factor=4, stride=2, style='pytorch')的avd_layer.stride == 2,验证 avg_down_stride 生效;- 输入
(2, 64, 56, 56)经 Bottleneck 后输出形状不变,验证残差结构正确。
2. 骨干网络整体前向(test_resnest_backbone):
- 不支持的深度(如
depth=18)会抛出KeyError,因为arch_settings仅支持 [50, 101, 152, 200]; - 以
ResNeSt(depth=50, radix=2, reduction_factor=4, out_indices=(0, 1, 2, 3))前向 224x224 输入,四个输出阶段的特征形状分别为[2, 256, 56, 56]、[2, 512, 28, 28]、[2, 1024, 14, 14]、[2, 2048, 7, 7]——每个阶段通道数 ×2、分辨率减半,验证了 Split-Attention 骨干在多尺度特征提取上的正确性。
# tests/test_models/test_backbones/test_resnest.py feat = model(imgs) assert feat[0].shape == torch.Size([2, 256, 56, 56]) assert feat[1].shape == torch.Size([2, 512, 28, 28]) assert feat[2].shape == torch.Size([2, 1024, 14, 14]) assert feat[3].shape == torch.Size([2, 2048, 7, 7])七、如何将 ResNeSt 扩展到自己的模型
由于ResNeSt已注册进 MMSegmentation 的模型注册表,任何以 ResNet 为 backbone 的分割配置都可以通过三行改动迁移到 ResNeSt:
model = dict( pretrained='open-mmlab://resnest101', backbone=dict( type='ResNeSt', stem_channels=128, radix=2, reduction_factor=4, avg_down_stride=True))注意事项:
- 如果使用自己训练的 ResNeSt 权重,可将
pretrained替换为本地权重文件路径,或通过init_cfg指定; - 想调整 Split-Attention 的强度,可修改
radix(分支数)与reduction_factor(注意力瓶颈缩减比);radix=1时 Split-Attention 退化为 SE 式 sigmoid 门控; - 显存紧张时可在 backbone 中开启
with_cp=True(checkpoint)以时间换显存; - 所有配置中的
stem_channels=128必须与预训练权重结构一致,否则会因权重 shape 不匹配而加载失败。
八、总结
ResNeSt 通过 Split-Attention 将通道注意力和多路径表示统一为一个模块化计算块,在 MMSegmentation 中作为分割骨干网络提供了从 FCN 到 DeepLabV3+ 的完整覆盖。本文从论文背景、RSoftmax/SplitAttentionConv2d/Bottleneck的源码实现、8 个仓库配置的复用模式、Cityscapes 与 ADE20K 的基准结果,到训练测试推理的完整实战流程,系统梳理了 ResNeSt 在 MMSegmentation 中的接入方式。你可以直接参考 configs/resnest/ 下的配置与 configs/resnest/metafile.yaml 中的模型清单,将 ResNeSt 应用到自己的分割任务中。
【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考