| 48 | |
| 49 | |
| 50 | class DownsampleBlock3d(nn.Module): |
| 51 | def __init__( |
| 52 | self, |
| 53 | in_channels: int, |
| 54 | out_channels: int, |
| 55 | mode: Literal["conv", "avgpool"] = "conv", |
| 56 | ): |
| 57 | assert mode in ["conv", "avgpool"], f"Invalid mode {mode}" |
| 58 | |
| 59 | super().__init__() |
| 60 | self.in_channels = in_channels |
| 61 | self.out_channels = out_channels |
| 62 | |
| 63 | if mode == "conv": |
| 64 | self.conv = nn.Conv3d(in_channels, out_channels, 2, stride=2) |
| 65 | elif mode == "avgpool": |
| 66 | assert in_channels == out_channels, "Pooling mode requires in_channels to be equal to out_channels" |
| 67 | |
| 68 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 69 | if hasattr(self, "conv"): |
| 70 | return self.conv(x) |
| 71 | else: |
| 72 | return F.avg_pool3d(x, 2) |
| 73 | |
| 74 | |
| 75 | class UpsampleBlock3d(nn.Module): |