(
self,
in_channels: int,
out_channels: int,
num_blocks: Tuple[int, ...],
block_out_channels: Tuple[int, ...],
act_fn: str,
)
| 839 | """ |
| 840 | |
| 841 | def __init__( |
| 842 | self, |
| 843 | in_channels: int, |
| 844 | out_channels: int, |
| 845 | num_blocks: Tuple[int, ...], |
| 846 | block_out_channels: Tuple[int, ...], |
| 847 | act_fn: str, |
| 848 | ): |
| 849 | super().__init__() |
| 850 | |
| 851 | layers = [] |
| 852 | for i, num_block in enumerate(num_blocks): |
| 853 | num_channels = block_out_channels[i] |
| 854 | |
| 855 | if i == 0: |
| 856 | layers.append(nn.Conv2d(in_channels, num_channels, kernel_size=3, padding=1)) |
| 857 | else: |
| 858 | layers.append( |
| 859 | nn.Conv2d( |
| 860 | num_channels, |
| 861 | num_channels, |
| 862 | kernel_size=3, |
| 863 | padding=1, |
| 864 | stride=2, |
| 865 | bias=False, |
| 866 | ) |
| 867 | ) |
| 868 | |
| 869 | for _ in range(num_block): |
| 870 | layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn)) |
| 871 | |
| 872 | layers.append(nn.Conv2d(block_out_channels[-1], out_channels, kernel_size=3, padding=1)) |
| 873 | |
| 874 | self.layers = nn.Sequential(*layers) |
| 875 | self.gradient_checkpointing = False |
| 876 | |
| 877 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 878 | r"""The forward method of the `EncoderTiny` class.""" |
nothing calls this directly
no test coverage detected