GRN (Global Response Normalization) layer
| 326 | |
| 327 | |
| 328 | class GRN(nn.Module): |
| 329 | """ GRN (Global Response Normalization) layer |
| 330 | """ |
| 331 | |
| 332 | def __init__(self, dim): |
| 333 | super().__init__() |
| 334 | self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) |
| 335 | self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) |
| 336 | |
| 337 | def forward(self, x): |
| 338 | Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) |
| 339 | Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) |
| 340 | return self.gamma * (x * Nx) + self.beta + x |
| 341 | |
| 342 | |
| 343 | class LayerNorm(nn.Module): |