| 44 | |
| 45 | class autoencoder_vgg1(nn.Module): # psnr 20.84 |
| 46 | def __init__(self): |
| 47 | super(autoencoder_vgg1, self).__init__() |
| 48 | self.encoder = nn.Sequential( # like the Composition layer you built |
| 49 | nn.Conv2d(3, 16, 3, stride=2, padding=1), |
| 50 | nn.ReLU(), |
| 51 | nn.Conv2d(16, 32, 3, stride=2, padding=1), |
| 52 | nn.ReLU(), |
| 53 | nn.Conv2d(32, 64, 7) |
| 54 | ) |
| 55 | self.decoder = nn.Sequential( |
| 56 | nn.ConvTranspose2d(64, 32, 7), |
| 57 | nn.ReLU(), |
| 58 | nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1), |
| 59 | nn.ReLU(), |
| 60 | nn.ConvTranspose2d(16, 3, 3, stride=2, padding=1, output_padding=1), |
| 61 | # nn.Sigmoid() # if img is 0 to 1 |
| 62 | nn.Tanh() # if img is -1 to 1 |
| 63 | ) |
| 64 | def forward(self, x): |
| 65 | encode = self.encoder(x) |
| 66 | decode = self.decoder(encode) |