| 10 | return floor((img_size + 2 - 3) / 2.) + 1 |
| 11 | |
| 12 | class EncoderBlock(nn.Module): |
| 13 | def __init__(self, |
| 14 | in_channels: int, |
| 15 | out_channels: int, |
| 16 | latent_dim: int, |
| 17 | img_size: int): |
| 18 | super(EncoderBlock, self).__init__() |
| 19 | |
| 20 | # Build Encoder |
| 21 | self.encoder = nn.Sequential( |
| 22 | nn.Conv2d(in_channels, |
| 23 | out_channels, |
| 24 | kernel_size=3, stride=2, padding=1), |
| 25 | nn.BatchNorm2d(out_channels), |
| 26 | nn.LeakyReLU()) |
| 27 | |
| 28 | out_size = conv_out_shape(img_size) |
| 29 | self.encoder_mu = nn.Linear(out_channels * out_size ** 2 , latent_dim) |
| 30 | self.encoder_var = nn.Linear(out_channels * out_size ** 2, latent_dim) |
| 31 | |
| 32 | def forward(self, input: Tensor) -> Tensor: |
| 33 | result = self.encoder(input) |
| 34 | h = torch.flatten(result, start_dim=1) |
| 35 | |
| 36 | # Split the result into mu and var components |
| 37 | # of the latent Gaussian distribution |
| 38 | mu = self.encoder_mu(h) |
| 39 | log_var = self.encoder_var(h) |
| 40 | |
| 41 | return [result, mu, log_var] |
| 42 | |
| 43 | class LadderBlock(nn.Module): |
| 44 | def __init__(self, |