| 8 | |
| 9 | class ConvAutoencoder(nn.Module): |
| 10 | def __init__(self): |
| 11 | super(ConvAutoencoder, self).__init__() |
| 12 | ## encoder layers ## |
| 13 | # conv layer (depth from 3 --> 16), 3x3 kernels |
| 14 | self.conv1 = nn.Conv2d(3, 16, 3, padding=1) |
| 15 | # conv layer (depth from 16 --> 4), 3x3 kernels |
| 16 | self.conv2 = nn.Conv2d(16, 4, 3, padding=1) |
| 17 | # pooling layer to reduce x-y dims by two; kernel and stride of 2 |
| 18 | self.pool = nn.MaxPool2d(2, 2) |
| 19 | |
| 20 | ## decoder layers ## |
| 21 | ## a kernel of 2 and a stride of 2 will increase the spatial dims by 2 |
| 22 | self.t_conv1 = nn.ConvTranspose2d(4, 16, 2, stride=2) |
| 23 | self.t_conv2 = nn.ConvTranspose2d(16, 3, 2, stride=2) |
| 24 | |
| 25 | def forward(self, x): |
| 26 | ## encode ## |