Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002
| 1423 | |
| 1424 | |
| 1425 | class SpatialNorm(nn.Module): |
| 1426 | """ |
| 1427 | Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002 |
| 1428 | """ |
| 1429 | |
| 1430 | def __init__( |
| 1431 | self, |
| 1432 | f_channels, |
| 1433 | zq_channels, |
| 1434 | ): |
| 1435 | super().__init__() |
| 1436 | self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True) |
| 1437 | self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 1438 | self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 1439 | |
| 1440 | def forward(self, f, zq): |
| 1441 | f_size = f.shape[-2:] |
| 1442 | zq = F.interpolate(zq, size=f_size, mode="nearest") |
| 1443 | norm_f = self.norm_layer(f) |
| 1444 | new_f = norm_f * self.conv_y(zq) + self.conv_b(zq) |
| 1445 | return new_f |