64 lines
2.4 KiB
Python
64 lines
2.4 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
"""
|
|
@Author : Peike Li
|
|
@Contact : peike.li@yahoo.com
|
|
@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 |