2019 CVPR 《Selective Kernel Networks》 PyTorch实现

import numpy as np
import torch
from torch import nn
from torch.nn import init
from collections import OrderedDict


# selective kernel attention
# 多个卷积核的通道注意力
# 方法出处 2019 CVPR 《Selective Kernel Networks》
class SKAttention(nn.Module):
    # 初始化层
    def __init__(self, channel=512, kernels=[1, 3, 5, 7], reduction=16, group=1, L=32):
        # 所有继承于nn.Module的模型都要写这句话
        super(SKAttention, self).__init__()
        # 隐藏层的缩减因子
        self.d = max(L, channel // reduction)
        # 不同大小的卷积核
        # nn.ModuleList类似于Python中的list
        # 只不过在PyTorch中
        # ModuleList中存储的是神经网络层
        self.convs = nn.ModuleList([])
        # 遍历不同大小的卷积核
        for k in kernels:
            # 添加不同大小卷积核的卷积层
            self.convs.append(
                # nn.Sequential可以理解为不同网络层的一个组合器
                # 这个OrderedDict的作用是可以为加入到Sequential的每一个神经网络层进行命名
                # 这样我们可以通过名字找到相应的神经网络层
                nn.Sequential(OrderedDict([
                    # Conv2d卷积层的输入是[N,C-in,H,W]
                    # 表示[批大小，输入的通道数，输入的高，输入的宽]
                    # group表示分组卷积的组数
                    # 以group=2为例
                    # 将输入的通道数除以2,也就是将输入的数据一分为2
                    # 然后对于每一部分进行卷积
                    # 最后将结果进行连接
                    # 卷积层输出的宽，高计算公式
                    # 输出的宽=(输入的宽+2*padding-kernel_size)/stride(卷积步幅，默认为1)+1
                    # 所以最后不同卷积核输出的特征图的宽和高都是相同的
                    # 这样后面融合的步骤才能将这些不同卷积核输出的特征图进行相加
                    ('conv', nn.Conv2d(channel, channel, kernel_size=k, padding=k // 2, groups=group)),
                    ('bn', nn.BatchNorm2d(channel)),
                    ('relu', nn.ReLU())
                ]))
            )
        # 公用的通道衰减层
        self.fc = nn.Linear(channel, self.d)
        # 不同卷积核的通道注意力得分
        self.fcs = nn.ModuleList([])
        for i in range(len(kernels)):
            self.fcs.append(nn.Linear(self.d, channel))
        self.softmax = nn.Softmax(dim=0)

    def forward(self, x):
        # 获取输入特征图的批大小，通道数
        bs, c, _, _ = x.size()
        # 不同卷积核的卷积结果
        conv_outs = []
        # 论文中所述的split操作
        # 利用不同大小的卷积核对输入的特征图进行卷积
        for conv in self.convs:
            conv_outs.append(conv(x))
        # stack表示张量的堆叠,会在指定的维度上创建新的维度
        # cat表示张量的拼接，在指定的维度上进行张量连接
        feats = torch.stack(conv_outs, 0)  # k,bs,channel,h,w

        # 论文中所述的fuse操作
        # 按照第一个维度相加张量
        U = sum(conv_outs)  # bs,c,h,w

        # 论文所述的压缩特诊图的操作reduction channel
        # 相当于按宽，按高相加求均值
        S = U.mean(-1).mean(-1)  # bs,c
        # 得到隐藏层值
        Z = self.fc(S)  # bs,d
        # 计算不同通道之间的注意力得分
        weights = []
        for fc in self.fcs:
            # 利用隐藏层输出不同注意力得分
            weight = fc(Z)
            weights.append(weight.view(bs, c, 1, 1))  # bs,channel
        attention_weughts = torch.stack(weights, 0)  # k,bs,channel,1,1
        # 将注意力得分变为权重
        attention_weughts = self.softmax(attention_weughts)  # k,bs,channel,1,1

        # 论文所述的不同通道注意力融合的操作
        # 将不同卷积核输出的特征图乘以相应的权重
        # 最后再将这些特征图进行求和
        V = (attention_weughts * feats).sum(0)
        return V


if __name__ == '__main__':
    # 输入，可以想象成特征图
    # 批大小50
    # 通道512
    # 宽，高7*7
    input = torch.randn(50, 512, 7, 7)
    se = SKAttention(channel=512, reduction=8)
    output = se(input)
    print(output.shape)