| 41 | """ |
| 42 | |
| 43 | def __init__(self, num_latent_dims, image_shape, max_num_filters): |
| 44 | super().__init__() |
| 45 | |
| 46 | # number of filters in the convolutional layers |
| 47 | num_filters_1 = max_num_filters // 4 |
| 48 | num_filters_2 = max_num_filters // 2 |
| 49 | num_filters_3 = max_num_filters |
| 50 | |
| 51 | # Output (BHWC): B x 32 x 32 x num_filters_1 |
| 52 | self.conv1 = nn.Conv2d(image_shape[-1], num_filters_1, 3, stride=2, padding=1) |
| 53 | # Output (BHWC): B x 16 x 16 x num_filters_2 |
| 54 | self.conv2 = nn.Conv2d(num_filters_1, num_filters_2, 3, stride=2, padding=1) |
| 55 | # Output (BHWC): B x 8 x 8 x num_filters_3 |
| 56 | self.conv3 = nn.Conv2d(num_filters_2, num_filters_3, 3, stride=2, padding=1) |
| 57 | |
| 58 | # Batch Normalization |
| 59 | self.bn1 = nn.BatchNorm(num_filters_1) |
| 60 | self.bn2 = nn.BatchNorm(num_filters_2) |
| 61 | self.bn3 = nn.BatchNorm(num_filters_3) |
| 62 | |
| 63 | # Divide the spatial dimensions by 8 because of the 3 strided convolutions |
| 64 | output_shape = [num_filters_3] + [ |
| 65 | dimension // 8 for dimension in image_shape[:-1] |
| 66 | ] |
| 67 | |
| 68 | flattened_dim = math.prod(output_shape) |
| 69 | |
| 70 | # Linear mappings to mean and standard deviation |
| 71 | self.proj_mu = nn.Linear(flattened_dim, num_latent_dims) |
| 72 | self.proj_log_var = nn.Linear(flattened_dim, num_latent_dims) |
| 73 | |
| 74 | def __call__(self, x): |
| 75 | x = nn.leaky_relu(self.bn1(self.conv1(x))) |