We are in the process of training an AE on CIFAR10 images. We used the following architecture:
class OurAE(nn.Module):
def __init__(self, in_channels, z_channels):
super(OurAE, self).__init__()
self.tot_diff = None
self.in_channels = in_channels
curr_channels = in_channels
encoder = []
channels = [3, 16, 64] + [z_channels]
for out_channels in channels:
encoder += [
nn.Conv2d(in_channels=curr_channels, out_channels=out_channels, kernel_size=3, padding=1, stride=2),
nn.ReLU()
]
curr_channels = out_channels
self.encoder = nn.Sequential(*encoder)
curr_channels = z_channels
decoder = []
channels = [64, 16, 3] + [in_channels]
for out_channels in channels:
decoder += [
nn.ConvTranspose2d(in_channels=curr_channels, out_channels=out_channels, kernel_size=4, padding=1, stride=2),
nn.ReLU()
]
curr_channels = out_channels
decoder = decoder[:-1] # removing the RELU layer
decoder.append(nn.Sigmoid())
self.decoder = nn.Sequential(*decoder)
def forward(self, x):
return self.decoder(self.encoder(x))
We are not sure why but we always get black and white images.
We tried to replace the Sigmoid with ReLU in the last layer but to no avail.
These are the loss function and the optimizers that we used:
optimizer = torch.optim.Adam(classifier.parameters(), lr=lr)
criterion = torch.nn.CrossEntropyLoss()
Here is an example of an input and output of the AE after training:
