Implements the decoder side of the Autoencoder.
| 157 | |
| 158 | |
| 159 | class Decoder(nn.Module): |
| 160 | """Implements the decoder side of the Autoencoder.""" |
| 161 | |
| 162 | def __init__( |
| 163 | self, |
| 164 | in_channels: int, |
| 165 | out_channels: int, |
| 166 | block_out_channels: List[int] = [64], |
| 167 | layers_per_block: int = 2, |
| 168 | resnet_groups: int = 32, |
| 169 | ): |
| 170 | super().__init__() |
| 171 | |
| 172 | self.conv_in = nn.Conv2d( |
| 173 | in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1 |
| 174 | ) |
| 175 | |
| 176 | self.mid_blocks = [ |
| 177 | ResnetBlock2D( |
| 178 | in_channels=block_out_channels[-1], |
| 179 | out_channels=block_out_channels[-1], |
| 180 | groups=resnet_groups, |
| 181 | ), |
| 182 | Attention(block_out_channels[-1], resnet_groups), |
| 183 | ResnetBlock2D( |
| 184 | in_channels=block_out_channels[-1], |
| 185 | out_channels=block_out_channels[-1], |
| 186 | groups=resnet_groups, |
| 187 | ), |
| 188 | ] |
| 189 | |
| 190 | channels = list(reversed(block_out_channels)) |
| 191 | channels = [channels[0]] + channels |
| 192 | self.up_blocks = [ |
| 193 | EncoderDecoderBlock2D( |
| 194 | in_channels, |
| 195 | out_channels, |
| 196 | num_layers=layers_per_block, |
| 197 | resnet_groups=resnet_groups, |
| 198 | add_downsample=False, |
| 199 | add_upsample=i < len(block_out_channels) - 1, |
| 200 | ) |
| 201 | for i, (in_channels, out_channels) in enumerate(zip(channels, channels[1:])) |
| 202 | ] |
| 203 | |
| 204 | self.conv_norm_out = nn.GroupNorm( |
| 205 | resnet_groups, block_out_channels[0], pytorch_compatible=True |
| 206 | ) |
| 207 | self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1) |
| 208 | |
| 209 | def __call__(self, x): |
| 210 | x = self.conv_in(x) |
| 211 | |
| 212 | x = self.mid_blocks[0](x) |
| 213 | x = self.mid_blocks[1](x) |
| 214 | x = self.mid_blocks[2](x) |
| 215 | |
| 216 | for l in self.up_blocks: |