【YOLOv4探讨 之八】（2）空间注意力机制SAM -- 利用Darknet YOLOv4在网络中添加注意力机制

在上一篇《【YOLOv4探讨 之七】利用Darknet YOLOv4在网络中添加注意力机制模块 系列之SE模块》 https://blog.youkuaiyun.com/qq_41736617/article/details/118424585中，我们介绍了SE模块的添加方法，这一篇我们在Darknet中增加了SAM模块。

基本概念

空间注意力机制使用SAM模块，在Darknet中，新添加的sam_layer层就是用于SAM模块，该层在darknet.h中的定义为sam. 其原理图如下：

其在网络中的部位仍然是RES残差模块中，首先对残差模块最后一个卷积模块输出分别求沿着通道方向的全局maxpool和全局avgpool，形成两个通道数为1的feature map,对两个feature map做containation，然后对这个2通道的输出做卷积，卷积完毕后使用Sigmoid激活函数确定空间平面上的权重，然后和残差模块最后一个卷积模块输出相乘。

该过程主要功能是提升目标定位效果，在空间上突出需要定位的目标打分权重。

配置实现

这里依然使用的是yolov3-tiny.cfg进行改造，添加RES和SAM模块需要在配置文件中增加####标注的内容：

......
......
batch_normalize=1
filters=256
size=3
stride=1
pad=1
activation=leaky

[maxpool]
size=2
stride=2


[convolutional]
batch_normalize=1
filters=128
size=1
stride=1
pad=1
activation=leaky

#########新增的配置内容#######
####先对RES模块增加做准备，通道数一般往小设计，后续还要通过route层做containation###

[route]
layers = -2

[convolutional]
batch_normalize=1
filters=128
size=1
stride=1
pad=1
activation=leaky

####两个RES模块######
[convolutional]
batch_normalize=1
filters=128
size=1
stride=1
pad=1
activation=leaky

[convolutional]
batch_normalize=1
filters=128
size=3
stride=1
pad=1
activation=leaky

[shortcut]
from=-3
activation=linear

[convolutional]
batch_normalize=1
filters=128
size=1
stride=1
pad=1
activation=leaky

[convolutional]
batch_normalize=1
filters=128
size=3
stride=1
pad=1
activation=leaky

###SAM模块###
#通道方向全局最大池化
[maxpool]
maxpool_depth = 1
out_channels = 1
#通道方向全局平均池化
[route]
layers = -2

[avgpool]
channelpool = 1
#对两个1*H*W的池化层做containation
[route]
layers = -1, -3
#对containation后的池化层进行卷积
[convolutional]
batch_normalize=1
filters=128
size=7
stride=1
pad=1
activation=logistic#做Sigmoid
#空间注意力加权
[sam]
from = -6
activation= linear

###SAM模块结束####

[shortcut]
from=-9
activation=linear
###RES模块结束####

[convolutional]
batch_normalize=1
filters=128
size=1
stride=1
pad=1
activation=leaky

[route]
layers = -1,-16

[convolutional]
batch_normalize=1
filters=256
size=1
stride=1
pad=1
activation=leaky
#####新增的配置内容结束#####
[convolutional]
batch_normalize=1
filters=512
size=3
stride=1
pad=1
activation=leaky

[maxpool]
size=2
stride=1

[convolutional]
batch_normalize=1
filters=1024
size=3
stride=1
pad=1
activation=leaky
......
......

以上配置文件中，对containation后的池化层进行卷积这个过程这里直接进行128通道的卷积。严格按照原理图，卷积后为$1\times H\times W$，因为需要使用sam_layer和128通道的输入层进行相乘，这里需要进行128次containation。这个过程也可使用如下配置片段进行代替

......
......
#对containation后的池化层进行卷积
[convolutional]
batch_normalize=1
filters=1
size=7
stride=1
pad=1
activation=logistic#做Sigmoid
#空间注意力加权
[route]
layers = -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,\
-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,\
-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,\
-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,\
-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1

[sam]
from = -7
activation= linear

###SAM模块结束####

[shortcut]
from=-10
activation=linear
###RES模块结束####

[convolutional]
batch_normalize=1
filters=128
size=1
stride=1
pad=1
activation=leaky

[route]
layers = -1,-17

[convolutional]
batch_normalize=1
filters=256
size=1
stride=1
pad=1
activation=leaky
#####新增的配置内容结束#####
......
......

但是遇到一个最大的问题，训练一段时间后，会导致

网上都说是显存溢出在进行此操作时，同时监控显存，发现显存占用并没有爆。
猜测有可能是因为128通道每个通道的内容相同，导致反响传播时候出现梯度爆炸，同时上面结果中也可看出loss = -nan。
但在之前的测试中也出现过持续的nan之类，并没有因为loss = -nan或loss = nan程序立马崩溃，可见主要问题还是内存出错。因此更大的可能是因为反复堆叠feature map，内存不连贯，指针偶尔出错有关。
为了实现128层的扩展，使用了128通道的卷积可以实现类似功能，但不会出现梯度爆炸等情况，这里就采用这种方式。

源码修改与分析

这里主要用到YOLOv4新增的sam_layer.c。
由于Darknet中的avgpool_layer.c中没有通道方向的全局平均池化，本人在Darknet的代码中进行了修改，主要涉及parser.c,avgpool_layer.c,avgpool_layer.h和avgpool_layer_kennels.cu.
废话不多说，放码过来。

sam_layer

• parser.c

//parse_avgpool可以看出在Darknet框架中cfg文件中需要配置的参数为from和activation
//from就是指定将当前的SAM权重map和哪个层的feature map相乘
//activation默认为linear，同时不支持SWISH或MISH
layer parse_sam(list *options, size_params params, network net)
{
   
   
    char *l = option_find(options, "from");
    int index = atoi(l);
    if (index < 0) index = params.index + index;

    int batch = params.batch;
    layer from = net.layers[index];

    layer s = make_sam_layer(batch, index, params.w, params.h, params.c, from.out_w, from.out_h, from.out_c);

    char *activation_s = option_find_str_quiet(options, "activation", "linear");
    ACTIVATION activation = get_activation(activation_s);
    s.activation = activation;
    if (activation == SWISH || activation == MISH) {
   
   
        printf(" [sam] layer doesn't support SWISH or MISH activations \n");
    }
    return s;
}

• sam_layer.c

void forward_sam_layer(const layer l, network_state state)
{
   
   
    //计算输出feature map的尺寸
    int size = l.batch * l.out_c * l.out_w * l.out_h;
    
    float *from_output = state.net.layers[l.index].output;

    int i;
    #pragma omp parallel for
    for (i = 0; i < size; ++i) {
   
   
        //将SAM模块输出map和需要处理的feature map点乘
        //注意，输出的size设置为多大，SAM模块输出map有多少层，feature map就会选取的有多少层。另一层含义是和feature map相乘的SAM模块输出map要保持与之相同的size.
        l.output[i] = state.input[i] * from_output[i];
    }

    activate_array(l.output, l.outputs*l.batch, l.activation);
}

void backward_sam_layer(const layer l, network_state state)
{
   
   
    gradient_array(l.output, l.outputs*l.batch, l.activation, l.delta);
    //axpy_cpu(l.outputs*l.batch, 1, l.delta, 1, state.delta, 1);
    //scale_cpu(l.batch, l.out_w, l.out_h, l.out_c, l.delta, l.w, l.h, l.c, state.net.layers[l.index].delta);

    int size = l.batch * l.out_c * l.out_w * l.out_h;
    //int channel_size = 1;
    float *from_output = state.net