48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
"""
|
|
@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 |