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
| 117 | return x |
| 118 | |
| 119 | class LayerNorm(nn.Module): |
| 120 | r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. |
| 121 | The ordering of the dimensions in the inputs. channels_last corresponds to inputs with |
| 122 | shape (batch_size, height, width, channels) while channels_first corresponds to inputs |
| 123 | with shape (batch_size, channels, height, width). |
| 124 | """ |
| 125 | def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): |
| 126 | super().__init__() |
| 127 | self.weight = nn.Parameter(torch.ones(normalized_shape)) |
| 128 | self.bias = nn.Parameter(torch.zeros(normalized_shape)) |
| 129 | self.eps = eps |
| 130 | self.data_format = data_format |
| 131 | if self.data_format not in ["channels_last", "channels_first"]: |
| 132 | raise NotImplementedError |
| 133 | self.normalized_shape = (normalized_shape, ) |
| 134 | |
| 135 | def forward(self, x): |
| 136 | if self.data_format == "channels_last": |
| 137 | return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) |
| 138 | elif self.data_format == "channels_first": |
| 139 | u = x.mean(1, keepdim=True) |
| 140 | s = (x - u).pow(2).mean(1, keepdim=True) |
| 141 | x = (x - u) / torch.sqrt(s + self.eps) |
| 142 | x = self.weight[:, None, None] * x + self.bias[:, None, None] |
| 143 | return x |
| 144 | |
| 145 | |
| 146 | model_urls = { |