| 348 | |
| 349 | |
| 350 | class SpatialNorm3D(nn.Module): |
| 351 | def __init__( |
| 352 | self, |
| 353 | f_channels, |
| 354 | zq_channels, |
| 355 | freeze_norm_layer=False, |
| 356 | add_conv=False, |
| 357 | pad_mode="constant", |
| 358 | gather=False, |
| 359 | **norm_layer_params, |
| 360 | ): |
| 361 | super().__init__() |
| 362 | if gather: |
| 363 | self.norm_layer = ContextParallelGroupNorm(num_channels=f_channels, **norm_layer_params) |
| 364 | else: |
| 365 | self.norm_layer = torch.nn.GroupNorm(num_channels=f_channels, **norm_layer_params) |
| 366 | # self.norm_layer = norm_layer(num_channels=f_channels, **norm_layer_params) |
| 367 | if freeze_norm_layer: |
| 368 | for p in self.norm_layer.parameters: |
| 369 | p.requires_grad = False |
| 370 | |
| 371 | self.add_conv = add_conv |
| 372 | if add_conv: |
| 373 | self.conv = ContextParallelCausalConv3d( |
| 374 | chan_in=zq_channels, |
| 375 | chan_out=zq_channels, |
| 376 | kernel_size=3, |
| 377 | ) |
| 378 | |
| 379 | self.conv_y = ContextParallelCausalConv3d( |
| 380 | chan_in=zq_channels, |
| 381 | chan_out=f_channels, |
| 382 | kernel_size=1, |
| 383 | ) |
| 384 | self.conv_b = ContextParallelCausalConv3d( |
| 385 | chan_in=zq_channels, |
| 386 | chan_out=f_channels, |
| 387 | kernel_size=1, |
| 388 | ) |
| 389 | |
| 390 | def forward(self, f, zq): |
| 391 | if f.shape[2] == 1 and not _USE_CP: |
| 392 | zq = torch.nn.functional.interpolate(zq, size=f.shape[-3:], mode="nearest") |
| 393 | elif get_context_parallel_rank() == 0: |
| 394 | f_first, f_rest = f[:, :, :1], f[:, :, 1:] |
| 395 | f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:] |
| 396 | zq_first, zq_rest = zq[:, :, :1], zq[:, :, 1:] |
| 397 | zq_first = torch.nn.functional.interpolate(zq_first, size=f_first_size, mode="nearest") |
| 398 | zq_rest = torch.nn.functional.interpolate(zq_rest, size=f_rest_size, mode="nearest") |
| 399 | zq = torch.cat([zq_first, zq_rest], dim=2) |
| 400 | else: |
| 401 | zq = torch.nn.functional.interpolate(zq, size=f.shape[-3:], mode="nearest") |
| 402 | |
| 403 | if self.add_conv: |
| 404 | zq = self.conv(zq) |
| 405 | |
| 406 | # f = conv_gather_from_context_parallel_region(f, dim=2, kernel_size=1) |
| 407 | norm_f = self.norm_layer(f) |