(self, num_latent_dims, image_shape, max_num_filters)
| 95 | """A convolutional decoder""" |
| 96 | |
| 97 | def __init__(self, num_latent_dims, image_shape, max_num_filters): |
| 98 | super().__init__() |
| 99 | self.num_latent_dims = num_latent_dims |
| 100 | num_img_channels = image_shape[-1] |
| 101 | self.max_num_filters = max_num_filters |
| 102 | |
| 103 | # decoder layers |
| 104 | num_filters_1 = max_num_filters |
| 105 | num_filters_2 = max_num_filters // 2 |
| 106 | num_filters_3 = max_num_filters // 4 |
| 107 | |
| 108 | # divide the last two dimensions by 8 because of the 3 upsampling convolutions |
| 109 | self.input_shape = [dimension // 8 for dimension in image_shape[:-1]] + [ |
| 110 | num_filters_1 |
| 111 | ] |
| 112 | flattened_dim = math.prod(self.input_shape) |
| 113 | |
| 114 | # Output: flattened_dim |
| 115 | self.lin1 = nn.Linear(num_latent_dims, flattened_dim) |
| 116 | # Output (BHWC): B x 16 x 16 x num_filters_2 |
| 117 | self.upconv1 = UpsamplingConv2d( |
| 118 | num_filters_1, num_filters_2, 3, stride=1, padding=1 |
| 119 | ) |
| 120 | # Output (BHWC): B x 32 x 32 x num_filters_1 |
| 121 | self.upconv2 = UpsamplingConv2d( |
| 122 | num_filters_2, num_filters_3, 3, stride=1, padding=1 |
| 123 | ) |
| 124 | # Output (BHWC): B x 64 x 64 x #img_channels |
| 125 | self.upconv3 = UpsamplingConv2d( |
| 126 | num_filters_3, num_img_channels, 3, stride=1, padding=1 |
| 127 | ) |
| 128 | |
| 129 | # Batch Normalizations |
| 130 | self.bn1 = nn.BatchNorm(num_filters_2) |
| 131 | self.bn2 = nn.BatchNorm(num_filters_3) |
| 132 | |
| 133 | def __call__(self, z): |
| 134 | x = self.lin1(z) |
nothing calls this directly
no test coverage detected