ResNeSt 模型分析和代码详解 (拆组和通道注意力ResNet)

ResNeSt: Split-Attention Networks模型的拆分注意力网络，最近特别火，主要是作为深度学习的backbone模型，ResNeSt在不同的图像任务中都有效提高了模型的预测精度。
因此今天分享下，最近两天学习的心得体会，参考资料如下：
ResNeSt: Split-Attention Networks
github官网
B站作者讲解
张航主页

安装使用

由于作者的主要想法就是搭建一个方便大家使用的Deep learning的backbone。所以使用起来还是相当方便。安装直接使用pip就可以安装ResNeSt模块

pip install git+https://github.com/zhanghang1989/ResNeSt

使用就更简单啦，直接import

# using ResNeSt-50 as an example
from resnest.torch import resnest50
net = resnest50(pretrained=True)

简单测试一下代码的运行效果，对博客中的猫星人进行预测

import resnest
from resnest.torch import resnest50 
import torch
import numpy as np
from torchvision import transforms
from PIL import Image


def image_trans(img_file):

    tf = transforms.Compose([
        lambda x:Image.open(x).convert('RGB'), # string path= > image data
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                std=[0.229, 0.224, 0.225])
    ])
    img = tf(img_file).unsqueeze(0) # add batch dim

    return img

img = image_trans('./Cat.jpg')
net = resnest50(pretrained=True)
net.eval()

# input_tensor = torch.randn(2, 3, 224, 224)
output_tensor = net(img)
# [-1, 1] =>[0, 1] per softmax and transform to the numpy array
predictions = torch.nn.functional.softmax(output_tensor, 1).detach().numpy()

# ImageNet Decode
from tensorflow.keras.applications.mobilenet_v2 import  decode_predictions
# print(output_tensor.shape)
print('Predicted:', decode_predictions(predictions, top=3)[0])

Predicted: [(‘n02124075’, ‘Egyptian_cat’, 0.3746694), (‘n02123045’, ‘tabby’, 0.30778167), (‘n02123159’, ‘tiger_cat’, 0.04715328)]

对比MobileNet V2的预测结果，发现前两种的预测概率有所提升

Predicted: [(‘n02124075’, ‘Egyptian_cat’, 0.329623), (‘n02123045’, ‘tabby’, 0.2984233), (‘n02123159’, ‘tiger_cat’, 0.08494468)]

ResNeSt创新点

看论文的个人感受觉：读完对split attention的模型特点也不是很清楚，图看起来也不是太好理解，个人觉得可能是因为文中涉及的cardinality和radix两个超参的说明有点抽象，示意图也不是特别形象，导致很难一下就明白作者想表达的操作流程。
我是在看了源代码和作者讲解视频后才真正理解了模型的精髓点。其实。作者一直强调的两个点，即split(multi-brach)和channel-attention。如果理解了这两点，就算基本掌握ResNeSt。

split(multi-brach)分组再拆分！！！

• Multi-brach能更够使模型具有diverse representation，即在同一层中多个卷积核分支可以分别提取特征使得网络提取的特征更为多样。具体模型实现，作者在两个层级上体现了了mult-brach的思想。如下图中的cardinality 系数K和radix系数R，将输入特征分成了G组。
  对于cardinality层级，作者使用了nn.Conv2d中的groups参数将输入特征分成Cardinality=K组。这个K组特征层和卷积核的操作使相对独立的操作。
  对于radix层级，作者先使用了一个3x3的conv2d将通道数增加到channels*radix后，再用torch.split分成将特征层分成radix组输入。

Channel attention 小组求权重再融合！！！

• Channel-attention能够提供捕捉feature correlation的网络机制，即通过引入软注意机制实现特征通道间的权重分配。这里的权重系数获得可参考下图，首先将radix组的特征求和，再取global average pooling得到与单个radix组相同维度的矢量。
  然后使用两组1x1的conv2d进行权重系数的分配，具体维度上第一组卷积的输出维度为max(in_channels*radix//reduction_factor, 32) ，这里的reduction_factor =4缩放系数用于减少参数量；第二组的卷积的输出维度channels*radix，保持了与radix的输入特征层维度的对应。
  为保证radix组间的特征层的split的权重独立分布，使用r-softmax对各radix组的权重分别计算softmax得到attens，最后将各组对应的特征层与atten系数相乘再求和。

如果理解起来还是有点困难，没关系接下来我们对着源代码仔细剖析下ResNeSt的模型细节

代码详解

强烈建议对比TORCH官方的TORCHVISION.MODELS.RESNET50框架来来理解resnest的代码。
首先看看直接调用的resnest.torch.resnest.py文件中的内容。

resnest.py

定义了包括resnest50, 101, 200, 269在内的四个函数，函数内通过调用同目录下的resnet.py中的ResNet类建立model对象。

def resnest50(pretrained=False, root='~/.encoding/models', **kwargs):
    model = ResNet(Bottleneck, [3, 4, 6, 3],
                   radix=2, groups=1, bottleneck_width=64,
                   deep_stem=True, stem_width=32, avg_down=True,
                   avd=True, avd_first=False, **kwargs)
    if pretrained:
        model.load_state_dict(torch.hub.load_state_dict_from_url(
            resnest_model_urls['resnest50'], progress=True, check_hash=True))
    return model

对于pretrained weight 文件的加载，程序通过查找文件名对应的url地址进行在线下载和权重加载过程

# 地址类似https://s3.us-west-1.wasabisys.com/resnest/torch/resnest50-528c19ca.pth
_url_format = 'https://s3.us-west-1.wasabisys.com/resnest/torch/{}-{}.pth'

# 建立一个{"resnest50":528c19ca,,,}的字典
_model_sha256 = {
   name: checksum for checksum, name in [
    ('528c19ca', 'resnest50'),
    ('22405ba7', 'resnest101'),
    ('75117900', 'resnest200'),
    ('0cc87c48', 'resnest269'),
    ]}

def short_hash(