|
|
|
from torch import nn
|
|
from torch.autograd import Function
|
|
from torch.nn.modules.utils import _pair
|
|
|
|
from ..utils import ext_loader
|
|
|
|
ext_module = ext_loader.load_ext('_ext',
|
|
['psamask_forward', 'psamask_backward'])
|
|
|
|
|
|
class PSAMaskFunction(Function):
|
|
|
|
@staticmethod
|
|
def symbolic(g, input, psa_type, mask_size):
|
|
return g.op(
|
|
'mmcv::MMCVPSAMask',
|
|
input,
|
|
psa_type_i=psa_type,
|
|
mask_size_i=mask_size)
|
|
|
|
@staticmethod
|
|
def forward(ctx, input, psa_type, mask_size):
|
|
ctx.psa_type = psa_type
|
|
ctx.mask_size = _pair(mask_size)
|
|
ctx.save_for_backward(input)
|
|
|
|
h_mask, w_mask = ctx.mask_size
|
|
batch_size, channels, h_feature, w_feature = input.size()
|
|
assert channels == h_mask * w_mask
|
|
output = input.new_zeros(
|
|
(batch_size, h_feature * w_feature, h_feature, w_feature))
|
|
|
|
ext_module.psamask_forward(
|
|
input,
|
|
output,
|
|
psa_type=psa_type,
|
|
num_=batch_size,
|
|
h_feature=h_feature,
|
|
w_feature=w_feature,
|
|
h_mask=h_mask,
|
|
w_mask=w_mask,
|
|
half_h_mask=(h_mask - 1) // 2,
|
|
half_w_mask=(w_mask - 1) // 2)
|
|
return output
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
input = ctx.saved_tensors[0]
|
|
psa_type = ctx.psa_type
|
|
h_mask, w_mask = ctx.mask_size
|
|
batch_size, channels, h_feature, w_feature = input.size()
|
|
grad_input = grad_output.new_zeros(
|
|
(batch_size, channels, h_feature, w_feature))
|
|
ext_module.psamask_backward(
|
|
grad_output,
|
|
grad_input,
|
|
psa_type=psa_type,
|
|
num_=batch_size,
|
|
h_feature=h_feature,
|
|
w_feature=w_feature,
|
|
h_mask=h_mask,
|
|
w_mask=w_mask,
|
|
half_h_mask=(h_mask - 1) // 2,
|
|
half_w_mask=(w_mask - 1) // 2)
|
|
return grad_input, None, None, None
|
|
|
|
|
|
psa_mask = PSAMaskFunction.apply
|
|
|
|
|
|
class PSAMask(nn.Module):
|
|
|
|
def __init__(self, psa_type, mask_size=None):
|
|
super(PSAMask, self).__init__()
|
|
assert psa_type in ['collect', 'distribute']
|
|
if psa_type == 'collect':
|
|
psa_type_enum = 0
|
|
else:
|
|
psa_type_enum = 1
|
|
self.psa_type_enum = psa_type_enum
|
|
self.mask_size = mask_size
|
|
self.psa_type = psa_type
|
|
|
|
def forward(self, input):
|
|
return psa_mask(input, self.psa_type_enum, self.mask_size)
|
|
|
|
def __repr__(self):
|
|
s = self.__class__.__name__
|
|
s += f'(psa_type={self.psa_type}, '
|
|
s += f'mask_size={self.mask_size})'
|
|
return s
|
|
|
