(
self,
f_channels,
zq_channels,
freeze_norm_layer=False,
add_conv=False,
pad_mode="constant",
gather=False,
**norm_layer_params,
)
| 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: |
nothing calls this directly
no test coverage detected