face_autoencoder.py

class FaceAutoencoder(nn.Module):
    def __init__(self, n_out_channels1=4, n_out_channels2=4, n_out_channels3=1, \
                kernel_size1=5, kernel_size2=5, kernel_size3=5):
        super(FaceAutoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(N_IN_CHANNELS, out_channels=n_out_channels1, kernel_size=kernel_size1, stride=2, padding=2), # [3,160,160] -> [4, 80, 80]
            nn.ReLU(),

            nn.Dropout(0.3),

            nn.Conv2d(in_channels=n_out_channels1, out_channels=n_out_channels2, kernel_size=kernel_size2, stride=2, padding=2), # [3,80,80] -> [4, 40, 40]
            nn.ReLU(),
            nn.MaxPool2d(3, stride=1, padding=0), # [4, 40, 40] -> [4, 38, 38]

            nn.Dropout(0.3),

            nn.Conv2d(in_channels=n_out_channels2, out_channels=n_out_channels3, kernel_size=kernel_size3, stride=1, padding=1), # [4, 36, 36] -> [1, 36, 36]
            nn.MaxPool2d(5, stride=1, padding=2),
            nn.ReLU(),
        )
        self.decoder = nn.Sequential(
            nn.Conv2d(n_out_channels3, n_out_channels2, kernel_size1, stride=1, padding=2), # [1, 36, 36] -> [4, 36, 36]
            nn.ReLU(),
            # nn.MaxPool2d(5, stride=3, padding=2),
            nn.Upsample(size=(40, 40)),

            nn.Dropout(0.3),

            nn.Conv2d(n_out_channels2, n_out_channels1, kernel_size2, stride=1, padding=2),
            nn.ReLU(),
            nn.Upsample(scale_factor=2),

            nn.Dropout(0.3),

            nn.Conv2d(n_out_channels1, N_IN_CHANNELS, kernel_size3, stride=1, padding=2),
            nn.ReLU(),
            nn.Upsample(scale_factor=2),
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

    def encode(self, x):
        return self.encoder(x)