Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| # -*- encoding: utf-8 -*- | |
| """ | |
| @Author : Peike Li | |
| @Contact : [email protected] | |
| @File : aspp.py | |
| @Time : 8/4/19 3:36 PM | |
| @Desc : | |
| @License : This source code is licensed under the license found in the | |
| LICENSE file in the root directory of this source tree. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| from torch.nn import functional as F | |
| from modules import InPlaceABNSync | |
| class ASPPModule(nn.Module): | |
| """ | |
| Reference: | |
| Chen, Liang-Chieh, et al. *"Rethinking Atrous Convolution for Semantic Image Segmentation."* | |
| """ | |
| def __init__(self, features, out_features=512, inner_features=256, dilations=(12, 24, 36)): | |
| super(ASPPModule, self).__init__() | |
| self.conv1 = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), | |
| nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1, | |
| bias=False), | |
| InPlaceABNSync(inner_features)) | |
| self.conv2 = nn.Sequential( | |
| nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1, bias=False), | |
| InPlaceABNSync(inner_features)) | |
| self.conv3 = nn.Sequential( | |
| nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=False), | |
| InPlaceABNSync(inner_features)) | |
| self.conv4 = nn.Sequential( | |
| nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[1], dilation=dilations[1], bias=False), | |
| InPlaceABNSync(inner_features)) | |
| self.conv5 = nn.Sequential( | |
| nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[2], dilation=dilations[2], bias=False), | |
| InPlaceABNSync(inner_features)) | |
| self.bottleneck = nn.Sequential( | |
| nn.Conv2d(inner_features * 5, out_features, kernel_size=1, padding=0, dilation=1, bias=False), | |
| InPlaceABNSync(out_features), | |
| nn.Dropout2d(0.1) | |
| ) | |
| def forward(self, x): | |
| _, _, h, w = x.size() | |
| feat1 = F.interpolate(self.conv1(x), size=(h, w), mode='bilinear', align_corners=True) | |
| feat2 = self.conv2(x) | |
| feat3 = self.conv3(x) | |
| feat4 = self.conv4(x) | |
| feat5 = self.conv5(x) | |
| out = torch.cat((feat1, feat2, feat3, feat4, feat5), 1) | |
| bottle = self.bottleneck(out) | |
| return bottle |