r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, width, channels) while channels_first corresponds to inputs with shape (batch_siz
| 17 | from core.data.transforms.post_transforms import flip_back, pose_pck_accuracy, transform_preds |
| 18 | |
| 19 | class LayerNorm(nn.Module): |
| 20 | r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. |
| 21 | The ordering of the dimensions in the inputs. channels_last corresponds to inputs with |
| 22 | shape (batch_size, height, width, channels) while channels_first corresponds to inputs |
| 23 | with shape (batch_size, channels, height, width). |
| 24 | """ |
| 25 | def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): |
| 26 | super().__init__() |
| 27 | self.weight = nn.Parameter(torch.ones(normalized_shape)) |
| 28 | self.bias = nn.Parameter(torch.zeros(normalized_shape)) |
| 29 | self.eps = eps |
| 30 | self.data_format = data_format |
| 31 | if self.data_format not in ["channels_last", "channels_first"]: |
| 32 | raise NotImplementedError |
| 33 | self.normalized_shape = (normalized_shape, ) |
| 34 | |
| 35 | def forward(self, x): |
| 36 | if self.data_format == "channels_last": |
| 37 | return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) |
| 38 | elif self.data_format == "channels_first": |
| 39 | u = x.mean(1, keepdim=True) |
| 40 | s = (x - u).pow(2).mean(1, keepdim=True) |
| 41 | x = (x - u) / torch.sqrt(s + self.eps) |
| 42 | x = self.weight[:, None, None] * x + self.bias[:, None, None] |
| 43 | return x |
| 44 | |
| 45 | def resize(input, |
| 46 | size=None, |
no outgoing calls
no test coverage detected