(
self,
in_channels: int,
out_channels: int,
num_blocks: tuple[int, ...],
block_out_channels: tuple[int, ...],
act_fn: str,
)
| 770 | """ |
| 771 | |
| 772 | def __init__( |
| 773 | self, |
| 774 | in_channels: int, |
| 775 | out_channels: int, |
| 776 | num_blocks: tuple[int, ...], |
| 777 | block_out_channels: tuple[int, ...], |
| 778 | act_fn: str, |
| 779 | ): |
| 780 | super().__init__() |
| 781 | |
| 782 | layers = [] |
| 783 | for i, num_block in enumerate(num_blocks): |
| 784 | num_channels = block_out_channels[i] |
| 785 | |
| 786 | if i == 0: |
| 787 | layers.append(nn.Conv2d(in_channels, num_channels, kernel_size=3, padding=1)) |
| 788 | else: |
| 789 | layers.append( |
| 790 | nn.Conv2d( |
| 791 | num_channels, |
| 792 | num_channels, |
| 793 | kernel_size=3, |
| 794 | padding=1, |
| 795 | stride=2, |
| 796 | bias=False, |
| 797 | ) |
| 798 | ) |
| 799 | |
| 800 | for _ in range(num_block): |
| 801 | layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn)) |
| 802 | |
| 803 | layers.append(nn.Conv2d(block_out_channels[-1], out_channels, kernel_size=3, padding=1)) |
| 804 | |
| 805 | self.layers = nn.Sequential(*layers) |
| 806 | self.gradient_checkpointing = False |
| 807 | |
| 808 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 809 | r"""The forward method of the `EncoderTiny` class.""" |
nothing calls this directly
no test coverage detected