|
- # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
- from torch import nn
- from torch.autograd import Function
- from torch.autograd.function import once_differentiable
-
- from tensormask import _C
-
-
- class _SwapAlign2Nat(Function):
- @staticmethod
- def forward(ctx, X, lambda_val, pad_val):
- ctx.lambda_val = lambda_val
- ctx.input_shape = X.size()
-
- Y = _C.swap_align2nat_forward(X, lambda_val, pad_val)
- return Y
-
- @staticmethod
- @once_differentiable
- def backward(ctx, gY):
- lambda_val = ctx.lambda_val
- bs, ch, h, w = ctx.input_shape
-
- gX = _C.swap_align2nat_backward(gY, lambda_val, bs, ch, h, w)
-
- return gX, None, None
-
-
- swap_align2nat = _SwapAlign2Nat.apply
-
-
- class SwapAlign2Nat(nn.Module):
- """
- The op `SwapAlign2Nat` described in https://arxiv.org/abs/1903.12174.
- Given an input tensor that predicts masks of shape (N, C=VxU, H, W),
- apply the op, it will return masks of shape (N, V'xU', H', W') where
- the unit lengths of (V, U) and (H, W) are swapped, and the mask representation
- is transformed from aligned to natural.
- Args:
- lambda_val (int): the relative unit length ratio between (V, U) and (H, W),
- as we always have larger unit lengths for (V, U) than (H, W),
- lambda_val is always >= 1.
- pad_val (float): padding value for the values falling outside of the input
- tensor, default set to -6 as sigmoid(-6) is ~0, indicating
- that is no masks outside of the tensor.
- """
-
- def __init__(self, lambda_val, pad_val=-6.):
- super(SwapAlign2Nat, self).__init__()
- self.lambda_val = lambda_val
- self.pad_val = pad_val
-
- def forward(self, X):
- return swap_align2nat(X, self.lambda_val, self.pad_val)
-
- def __repr__(self):
- tmpstr = self.__class__.__name__ + "("
- tmpstr += "lambda_val=" + str(self.lambda_val)
- tmpstr += ", pad_val=" + str(self.pad_val)
- tmpstr += ")"
- return tmpstr
|
