def __init__(self, cfg, **kwargs):
super(PSANet, self).__init__(cfg, **kwargs)
align_corners, norm_cfg, act_cfg = self.align_corners, self.norm_cfg, self.act_cfg
# build psa
psa_cfg = cfg['psa']
assert psa_cfg['type'] in ['collect', 'distribute', 'bi-direction']
mask_h, mask_w = psa_cfg['mask_size']
if 'normalization_factor' not in psa_cfg:
psa_cfg['normalization_factor'] = mask_h * mask_w
self.reduce = nn.Sequential(
nn.Conv2d(psa_cfg['in_channels'], psa_cfg['out_channels'], kernel_size=1, stride=1, padding=0, bias=False),
BuildNormalization(norm_cfg['type'], (psa_cfg['out_channels'], norm_cfg['opts'])),
BuildActivation(act_cfg['type'], **act_cfg['opts']),
)
self.attention = nn.Sequential(
nn.Conv2d(psa_cfg['out_channels'], psa_cfg['out_channels'], kernel_size=1, stride=1, padding=0, bias=False),
BuildNormalization(norm_cfg['type'], (psa_cfg['out_channels'], norm_cfg['opts'])),
BuildActivation(act_cfg['type'], **act_cfg['opts']),
nn.Conv2d(psa_cfg['out_channels'], mask_h * mask_w, kernel_size=1, stride=1, padding=0, bias=False),
)
if psa_cfg['type'] == 'bi-direction':
self.reduce_p = nn.Sequential(
nn.Conv2d(psa_cfg['in_channels'], psa_cfg['out_channels'], kernel_size=1, stride=1, padding=0, bias=False),
BuildNormalization(norm_cfg['type'], (psa_cfg['out_channels'], norm_cfg['opts'])),
BuildActivation(act_cfg['type'], **act_cfg['opts']),
)
self.attention_p = nn.Sequential(
nn.Conv2d(psa_cfg['out_channels'], psa_cfg['out_channels'], kernel_size=1, stride=1, padding=0, bias=False),
BuildNormalization(norm_cfg['type'], (psa_cfg['out_channels'], norm_cfg['opts'])),
BuildActivation(act_cfg['type'], **act_cfg['opts']),
nn.Conv2d(psa_cfg['out_channels'], mask_h * mask_w, kernel_size=1, stride=1, padding=0, bias=False),
)
if not psa_cfg['compact']:
self.psamask_collect = PSAMask('collect', psa_cfg['mask_size'])
self.psamask_distribute = PSAMask('distribute', psa_cfg['mask_size'])
else:
if not psa_cfg['compact']:
self.psamask = PSAMask(psa_cfg['type'], psa_cfg['mask_size'])
self.proj = nn.Sequential(
nn.Conv2d(psa_cfg['out_channels'] * (2 if psa_cfg['type'] == 'bi-direction' else 1), psa_cfg['in_channels'], kernel_size=1, stride=1, padding=1, bias=False),
BuildNormalization(norm_cfg['type'], (psa_cfg['in_channels'], norm_cfg['opts'])),
BuildActivation(act_cfg['type'], **act_cfg['opts']),
)
# build decoder
decoder_cfg = cfg['decoder']
self.decoder = nn.Sequential(
nn.Conv2d(decoder_cfg['in_channels'], decoder_cfg['out_channels'], kernel_size=3, stride=1, padding=1, bias=False),
BuildNormalization(norm_cfg['type'], (decoder_cfg['out_channels'], norm_cfg['opts'])),
BuildActivation(act_cfg['type'], **act_cfg['opts']),
nn.Dropout2d(decoder_cfg['dropout']),
nn.Conv2d(decoder_cfg['out_channels'], cfg['num_classes'], kernel_size=1, stride=1, padding=0)
)
# build auxiliary decoder
self.setauxiliarydecoder(cfg['auxiliary'])
# freeze normalization layer if necessary
if cfg.get('is_freeze_norm', False): self.freezenormalization()
def test_psa_mask_distribute(self):
if not torch.cuda.is_available():
return
from mmcv.ops import PSAMask
test_loss = Loss()
input = np.fromfile(
'tests/data/for_psa_mask/psa_input.bin', dtype=np.float32)
output_distribute = np.fromfile(
'tests/data/for_psa_mask/psa_output_distribute.bin',
dtype=np.float32)
input = input.reshape((4, 16, 8, 8))
output_distribute = output_distribute.reshape((4, 64, 8, 8))
label = torch.ones((4, 64, 8, 8))
input = torch.FloatTensor(input)
input.requires_grad = True
psamask_distribute = PSAMask('distribute', (4, 4))
# test distribute cpu
test_output = psamask_distribute(input)
loss = test_loss(test_output, label)
loss.backward()
test_output = test_output.detach().numpy()
assert np.allclose(test_output, output_distribute)
assert test_output.shape == output_distribute.shape
psamask_distribute.cuda()
input = input.cuda()
label = label.cuda()
# test distribute cuda
test_output = psamask_distribute(input)
loss = test_loss(test_output, label)
loss.backward()
test_output = test_output.detach().cpu().numpy()
assert np.allclose(test_output, output_distribute)
assert test_output.shape == output_distribute.shape
def __init__(self,
mask_size,
psa_type='bi-direction',
compact=False,
shrink_factor=2,
normalization_factor=1.0,
psa_softmax=True,
**kwargs):
if PSAMask is None:
raise RuntimeError('Please install mmcv-full for PSAMask ops')
super(PSAHead, self).__init__(**kwargs)
assert psa_type in ['collect', 'distribute', 'bi-direction']
self.psa_type = psa_type
self.compact = compact
self.shrink_factor = shrink_factor
self.mask_size = mask_size
mask_h, mask_w = mask_size
self.psa_softmax = psa_softmax
if normalization_factor is None:
normalization_factor = mask_h * mask_w
self.normalization_factor = normalization_factor
self.reduce = ConvModule(self.in_channels,
self.channels,
kernel_size=1,
conv_cfg=self.conv_cfg,
norm_cfg=self.norm_cfg,
act_cfg=self.act_cfg)
self.attention = nn.Sequential(
ConvModule(self.channels,
self.channels,
kernel_size=1,
conv_cfg=self.conv_cfg,
norm_cfg=self.norm_cfg,
act_cfg=self.act_cfg),
nn.Conv2d(self.channels,
mask_h * mask_w,
kernel_size=1,
bias=False))
if psa_type == 'bi-direction':
self.reduce_p = ConvModule(self.in_channels,
self.channels,
kernel_size=1,
conv_cfg=self.conv_cfg,
norm_cfg=self.norm_cfg,
act_cfg=self.act_cfg)
self.attention_p = nn.Sequential(
ConvModule(self.channels,
self.channels,
kernel_size=1,
conv_cfg=self.conv_cfg,
norm_cfg=self.norm_cfg,
act_cfg=self.act_cfg),
nn.Conv2d(self.channels,
mask_h * mask_w,
kernel_size=1,
bias=False))
self.psamask_collect = PSAMask('collect', mask_size)
self.psamask_distribute = PSAMask('distribute', mask_size)
else:
self.psamask = PSAMask(psa_type, mask_size)
self.proj = ConvModule(self.channels *
(2 if psa_type == 'bi-direction' else 1),
self.in_channels,
kernel_size=1,
padding=1,
conv_cfg=self.conv_cfg,
norm_cfg=self.norm_cfg,
act_cfg=self.act_cfg)
self.bottleneck = ConvModule(self.in_channels * 2,
self.channels,
kernel_size=3,
padding=1,
conv_cfg=self.conv_cfg,
norm_cfg=self.norm_cfg,
act_cfg=self.act_cfg)