| 7 | from typing import List |
| 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 ## |
| 27 | # add hidden layers with relu activation function |
| 28 | # and maxpooling after |
| 29 | x = F.relu(self.conv1(x)) |
| 30 | x = self.pool(x) |
| 31 | # add second hidden layer |
| 32 | x = F.relu(self.conv2(x)) |
| 33 | x = self.pool(x) # compressed representation |
| 34 | |
| 35 | ## decode ## |
| 36 | # add transpose conv layers, with relu activation function |
| 37 | x = F.relu(self.t_conv1(x)) |
| 38 | # # output layer (with tanh for scaling from -1 to 1) |
| 39 | x = F.tanh(self.t_conv2(x)) |
| 40 | # output layer (with tanh for scaling from 0 to 1) |
| 41 | # x = F.sigmoid(self.t_conv2(x)) |
| 42 | |
| 43 | return x |
| 44 | |
| 45 | class autoencoder_vgg1(nn.Module): # psnr 20.84 |
| 46 | def __init__(self): |
nothing calls this directly
no outgoing calls
no test coverage detected