| 73 | |
| 74 | |
| 75 | class UpsampleBlock3d(nn.Module): |
| 76 | def __init__( |
| 77 | self, |
| 78 | in_channels: int, |
| 79 | out_channels: int, |
| 80 | mode: Literal["conv", "nearest"] = "conv", |
| 81 | ): |
| 82 | assert mode in ["conv", "nearest"], f"Invalid mode {mode}" |
| 83 | |
| 84 | super().__init__() |
| 85 | self.in_channels = in_channels |
| 86 | self.out_channels = out_channels |
| 87 | |
| 88 | if mode == "conv": |
| 89 | self.conv = nn.Conv3d(in_channels, out_channels*8, 3, padding=1) |
| 90 | elif mode == "nearest": |
| 91 | assert in_channels == out_channels, "Nearest mode requires in_channels to be equal to out_channels" |
| 92 | |
| 93 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 94 | if hasattr(self, "conv"): |
| 95 | x = self.conv(x) |
| 96 | return pixel_shuffle_3d(x, 2) |
| 97 | else: |
| 98 | return F.interpolate(x, scale_factor=2, mode="nearest") |
| 99 | |
| 100 | |
| 101 | class SparseStructureEncoder(nn.Module): |