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