【語義分割】PSPNet中PSP模塊的pytorch實現

github地址：https://github.com/Lextal/pspnet-pytorch/blob/master/pspnet.py

PSP模塊示意圖如下

 

代碼如下

class PSPModule(nn.Module):
    def __init__(self, features, out_features=1024, sizes=(1, 2, 3, 6)):
        super().__init__()
        self.stages = []
        self.stages = nn.ModuleList([self._make_stage(features, size) for size in sizes])
        self.bottleneck = nn.Conv2d(features * (len(sizes) + 1), out_features, kernel_size=1)
        self.relu = nn.ReLU()

    def _make_stage(self, features, size):
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        conv = nn.Conv2d(features, features, kernel_size=1, bias=False)
        return nn.Sequential(prior, conv)

    def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        priors = [F.upsample(input=stage(feats), size=(h, w), mode='bilinear') for stage in self.stages] + [feats]
        bottle = self.bottleneck(torch.cat(priors, 1))
        return self.relu(bottle)

此外，我基於自己的工作稍加修改，也給出一個3D版本。改動有幾處，一是3d卷積和池化，二是上采樣由雙線性插值切換為trilinear，不知是否翻譯為三線性插值，三是我對池化部分輸出尺寸的修改，上采樣到輸入的一半，同時與普通池化相結合，不過，這樣有沒有效果，我還沒試過

class PSPModule(nn.Module):
    def __init__(self, features, sizes=(1, 2, 3, 6)):
        super(PSPModule, self).__init__()
        self.stages = []
        self.stages = nn.ModuleList([self._make_stage(features, size) for size in sizes])
        self.bottleneck = nn.Conv3d(features * (1 + 1), features, kernel_size=1)
        self.relu = nn.ReLU()

    def _make_stage(self, features, size):
        prior = nn.AdaptiveAvgPool3d(output_size=(size, size, size))
        conv = nn.Conv3d(features, features / 4, kernel_size=1, bias=False)
        return nn.Sequential(prior, conv)

    def forward(self, x, maxpool_x):
        h, w, l = x.size(2), x.size(3), x.size(4)
        priors = [F.upsample(input=stage(x), size=(h / 2, w / 2, l/2), mode='trilinear') for stage in self.stages] + [maxpool_x]
        bottle = self.bottleneck(torch.cat(priors, 1))
        print(bottle.size())
        return self.relu(bottle)