Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| # -*- encoding: utf-8 -*- | |
| """ | |
| @Author : Peike Li | |
| @Contact : [email protected] | |
| @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 |