| |
| 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 |
|
|
