| | |
| | |
| |
|
| | """ |
| | @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 |
