61 lines
2.1 KiB
Python
61 lines
2.1 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import torch
|
|
import torch.nn as nn
|
|
from mmcv.cnn import ConvModule
|
|
|
|
from mmseg.ops import resize
|
|
from mmseg.models.decode_heads.decode_head import BaseDecodeHead
|
|
|
|
|
|
class PPM(nn.ModuleList):
|
|
"""Pooling Pyramid Module used in PSPNet.
|
|
|
|
Args:
|
|
pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid
|
|
Module.
|
|
in_channels (int): Input channels.
|
|
channels (int): Channels after modules, before conv_seg.
|
|
conv_cfg (dict|None): Config of conv layers.
|
|
norm_cfg (dict|None): Config of norm layers.
|
|
act_cfg (dict): Config of activation layers.
|
|
align_corners (bool): align_corners argument of F.interpolate.
|
|
"""
|
|
|
|
def __init__(self, pool_scales, in_channels, channels, conv_cfg, norm_cfg,
|
|
act_cfg, align_corners, **kwargs):
|
|
super(PPM, self).__init__()
|
|
self.pool_scales = pool_scales
|
|
self.align_corners = align_corners
|
|
self.in_channels = in_channels
|
|
self.channels = channels
|
|
self.conv_cfg = conv_cfg
|
|
self.norm_cfg = norm_cfg
|
|
self.act_cfg = act_cfg
|
|
for pool_scale in pool_scales:
|
|
self.append(
|
|
nn.Sequential(
|
|
nn.AdaptiveAvgPool2d(pool_scale),
|
|
ConvModule(
|
|
self.in_channels,
|
|
self.channels,
|
|
1,
|
|
conv_cfg=self.conv_cfg,
|
|
norm_cfg=self.norm_cfg,
|
|
act_cfg=self.act_cfg,
|
|
**kwargs)))
|
|
|
|
def forward(self, x):
|
|
"""Forward function."""
|
|
ppm_outs = []
|
|
for ppm in self:
|
|
ppm_out = ppm(x)
|
|
ppm_out = ppm_out.to(torch.float32)
|
|
upsampled_ppm_out = resize(
|
|
ppm_out,
|
|
size=x.size()[2:],
|
|
mode='bilinear',
|
|
align_corners=self.align_corners)
|
|
upsampled_ppm_out = upsampled_ppm_out.to(torch.bfloat16)
|
|
ppm_outs.append(upsampled_ppm_out)
|
|
return ppm_outs
|