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