ASPP模块

空洞空间卷积池化金字塔(atrous spatial pyramid pooling (ASPP))对所给定的输入以不同采样率的空洞卷积并行采样。

几个步骤

1. 1*1卷积

2. 3*3卷积( padding=6, dilation=6 )

3. 3*3 卷积( padding=12, dilation=12 )

4. 3*3 卷积( padding=18, dilation=18 )

5. 池化操作

  a.  全局平均池化 ----> 输出尺寸(1,1)

  b.  1*1卷积 改变通道

  c.  上采样 恢复成输入的尺寸

6.特征融合，此时输入通道为原始输入通道的5倍

代码如下：

from torch import nn
import torch
import torch.nn.functional as F
class ASPP(nn.Module):
    def __init__(self, in_channel):
        depth = in_channel
        super(ASPP,self).__init__()
        self.mean = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv2d(in_channel, depth, 1, 1)
        self.atrous_block1 = nn.Conv2d(in_channel, depth, 1, 1)
        self.atrous_block6 = nn.Conv2d(in_channel, depth, 3, 1, padding=6, dilation=6)
        self.atrous_block12 = nn.Conv2d(in_channel, depth, 3, 1, padding=12, dilation=12)
        self.atrous_block18 = nn.Conv2d(in_channel, depth, 3, 1, padding=18, dilation=18)
        self.conv_1x1_output = nn.Conv2d(depth * 5, depth, 1, 1)

    def forward(self, x):
        size = x.shape[2:]

        image_features = self.mean(x)
        image_features = self.conv(image_features)
        image_features = F.interpolate(image_features, size=size, mode='bilinear')

        atrous_block1 = self.atrous_block1(x)
        atrous_block6 = self.atrous_block6(x)
        atrous_block12 = self.atrous_block12(x)
        atrous_block18 = self.atrous_block18(x)

        cat = torch.cat([image_features, atrous_block1, atrous_block6,
                         atrous_block12, atrous_block18], dim=1)
        net = self.conv_1x1_output(cat)
        return net
aspp = ASPP(256)
out = torch.rand(2,256,13,13)
print(aspp(out).shape)