(
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,
)
| 917 | """ |
| 918 | |
| 919 | def __init__( |
| 920 | self, |
| 921 | in_channels: int, |
| 922 | out_channels: int, |
| 923 | num_blocks: Tuple[int, ...], |
| 924 | block_out_channels: Tuple[int, ...], |
| 925 | upsampling_scaling_factor: int, |
| 926 | act_fn: str, |
| 927 | upsample_fn: str, |
| 928 | ): |
| 929 | super().__init__() |
| 930 | |
| 931 | layers = [ |
| 932 | nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=1), |
| 933 | get_activation(act_fn), |
| 934 | ] |
| 935 | |
| 936 | for i, num_block in enumerate(num_blocks): |
| 937 | is_final_block = i == (len(num_blocks) - 1) |
| 938 | num_channels = block_out_channels[i] |
| 939 | |
| 940 | for _ in range(num_block): |
| 941 | layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn)) |
| 942 | |
| 943 | if not is_final_block: |
| 944 | layers.append(nn.Upsample(scale_factor=upsampling_scaling_factor, mode=upsample_fn)) |
| 945 | |
| 946 | conv_out_channel = num_channels if not is_final_block else out_channels |
| 947 | layers.append( |
| 948 | nn.Conv2d( |
| 949 | num_channels, |
| 950 | conv_out_channel, |
| 951 | kernel_size=3, |
| 952 | padding=1, |
| 953 | bias=is_final_block, |
| 954 | ) |
| 955 | ) |
| 956 | |
| 957 | self.layers = nn.Sequential(*layers) |
| 958 | self.gradient_checkpointing = False |
| 959 | |
| 960 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 961 | r"""The forward method of the `DecoderTiny` class.""" |
nothing calls this directly
no test coverage detected