(
self,
f_channels,
zq_channels,
freeze_norm_layer=False,
add_conv=False,
pad_mode="constant",
gather=False,
**norm_layer_params,
)
| 450 | |
| 451 | class SpatialNorm3D(nn.Module): |
| 452 | def __init__( |
| 453 | self, |
| 454 | f_channels, |
| 455 | zq_channels, |
| 456 | freeze_norm_layer=False, |
| 457 | add_conv=False, |
| 458 | pad_mode="constant", |
| 459 | gather=False, |
| 460 | **norm_layer_params, |
| 461 | ): |
| 462 | super().__init__() |
| 463 | if gather: |
| 464 | self.norm_layer = ContextParallelGroupNorm(num_channels=f_channels, **norm_layer_params) |
| 465 | else: |
| 466 | self.norm_layer = torch.nn.GroupNorm(num_channels=f_channels, **norm_layer_params) |
| 467 | # self.norm_layer = norm_layer(num_channels=f_channels, **norm_layer_params) |
| 468 | if freeze_norm_layer: |
| 469 | for p in self.norm_layer.parameters: |
| 470 | p.requires_grad = False |
| 471 | |
| 472 | self.add_conv = add_conv |
| 473 | if add_conv: |
| 474 | self.conv = ContextParallelCausalConv3d( |
| 475 | chan_in=zq_channels, |
| 476 | chan_out=zq_channels, |
| 477 | kernel_size=3, |
| 478 | ) |
| 479 | |
| 480 | self.conv_y = ContextParallelCausalConv3d( |
| 481 | chan_in=zq_channels, |
| 482 | chan_out=f_channels, |
| 483 | kernel_size=1, |
| 484 | ) |
| 485 | self.conv_b = ContextParallelCausalConv3d( |
| 486 | chan_in=zq_channels, |
| 487 | chan_out=f_channels, |
| 488 | kernel_size=1, |
| 489 | ) |
| 490 | |
| 491 | def forward(self, f, zq, clear_fake_cp_cache=True): |
| 492 | if f.shape[2] > 1 and f.shape[2] % 2 == 1: |
nothing calls this directly
no test coverage detected