ASPP（Atrous Spatial Pyramid Pooling）代码复现

import torch
import torch.nn as nn
import torch.nn.functional as F


class ASPP(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ASPP, self).__init__()
        dilations = [1, 6, 12, 18]
        self.layer1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1,
                      padding=0, dilation=dilations[0], bias=False),
            nn.BatchNorm2d(out_channels), nn.ReLU()
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1,
                      padding=dilations[1], dilation=dilations[1], bias=False),
            nn.BatchNorm2d(out_channels), nn.ReLU()
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1,
                      padding=dilations[2], dilation=dilations[2], bias=False),
            nn.BatchNorm2d(out_channels), nn.ReLU()
        )
        self.layer4 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1,
                      padding=dilations[3], dilation=dilations[3], bias=False),
            nn.BatchNorm2d(out_channels), nn.ReLU()
        )
        self.layer5 = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Conv2d(2048, 256, 1, stride=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU()
        )
        self.conv = nn.Conv2d(1280,256,1,bias=False)
        self.batch_nor = nn.BatchNorm2d(256)

    def forward(self,x):
        x1 = self.layer1(x)
        x2 = self.layer2(x)
        x3 = self.layer3(x)
        x4 = self.layer4(x)
        x5 = self.layer5(x)
        x5 = F.interpolate(x5,size=x1.size()[2:],mode='nilinear',align_corners=True)
        x = torch.cat((x1,x2,x3,x4,x5),dim=1)
        return x