| 47 | |
| 48 | |
| 49 | class ASPP(nn.Module): |
| 50 | def __init__(self, inplanes, mid_channels=256, BatchNorm=nn.BatchNorm2d): |
| 51 | super(ASPP, self).__init__() |
| 52 | dilations = [1, 6, 12, 18] |
| 53 | |
| 54 | self.aspp1 = _ASPPModule(inplanes, |
| 55 | mid_channels, |
| 56 | 1, |
| 57 | padding=0, |
| 58 | dilation=dilations[0], |
| 59 | BatchNorm=BatchNorm) |
| 60 | self.aspp2 = _ASPPModule(inplanes, |
| 61 | mid_channels, |
| 62 | 3, |
| 63 | padding=dilations[1], |
| 64 | dilation=dilations[1], |
| 65 | BatchNorm=BatchNorm) |
| 66 | self.aspp3 = _ASPPModule(inplanes, |
| 67 | mid_channels, |
| 68 | 3, |
| 69 | padding=dilations[2], |
| 70 | dilation=dilations[2], |
| 71 | BatchNorm=BatchNorm) |
| 72 | self.aspp4 = _ASPPModule(inplanes, |
| 73 | mid_channels, |
| 74 | 3, |
| 75 | padding=dilations[3], |
| 76 | dilation=dilations[3], |
| 77 | BatchNorm=BatchNorm) |
| 78 | |
| 79 | self.global_avg_pool = nn.Sequential( |
| 80 | nn.AdaptiveAvgPool2d((1, 1)), |
| 81 | nn.Conv2d(inplanes, mid_channels, 1, stride=1, bias=False), |
| 82 | BatchNorm(mid_channels), |
| 83 | nn.ReLU(), |
| 84 | ) |
| 85 | self.conv1 = nn.Conv2d(int(mid_channels * 5), |
| 86 | mid_channels, |
| 87 | 1, |
| 88 | bias=False) |
| 89 | self.bn1 = BatchNorm(mid_channels) |
| 90 | self.relu = nn.ReLU() |
| 91 | self.dropout = nn.Dropout(0.5) |
| 92 | self._init_weight() |
| 93 | |
| 94 | def forward(self, x): |
| 95 | x1 = self.aspp1(x) |
| 96 | x2 = self.aspp2(x) |
| 97 | x3 = self.aspp3(x) |
| 98 | x4 = self.aspp4(x) |
| 99 | x5 = self.global_avg_pool(x) |
| 100 | x5 = F.interpolate(x5, |
| 101 | size=x4.size()[2:], |
| 102 | mode='bilinear', |
| 103 | align_corners=True) |
| 104 | x = torch.cat((x1, x2, x3, x4, x5), dim=1) |
| 105 | |
| 106 | x = self.conv1(x) |