| |
| |
|
|
| """ |
| @Author : Peike Li |
| @Contact : peike.li@yahoo.com |
| @File : psp.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 PSPModule(nn.Module): |
| """ |
| Reference: |
| Zhao, Hengshuang, et al. *"Pyramid scene parsing network."* |
| """ |
| def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)): |
| super(PSPModule, self).__init__() |
|
|
| self.stages = [] |
| self.stages = nn.ModuleList([self._make_stage(features, out_features, size) for size in sizes]) |
| self.bottleneck = nn.Sequential( |
| nn.Conv2d(features + len(sizes) * out_features, out_features, kernel_size=3, padding=1, dilation=1, |
| bias=False), |
| InPlaceABNSync(out_features), |
| ) |
|
|
| def _make_stage(self, features, out_features, size): |
| prior = nn.AdaptiveAvgPool2d(output_size=(size, size)) |
| conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False) |
| bn = InPlaceABNSync(out_features) |
| return nn.Sequential(prior, conv, bn) |
|
|
| def forward(self, feats): |
| h, w = feats.size(2), feats.size(3) |
| priors = [F.interpolate(input=stage(feats), size=(h, w), mode='bilinear', align_corners=True) for stage in |
| self.stages] + [feats] |
| bottle = self.bottleneck(torch.cat(priors, 1)) |
| return bottle |
