Автоэнкодер выводит черно-белые изображения из цветных ⇐ Python

[Anonymous]

Мы находимся в процессе обучения AE на изображениях CIFAR10. Мы использовали следующую архитектуру:

Код: Выделить всё

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))

Мы не знаем почему, но мы всегда получаем черно-белые изображения.
Мы пытались заменить сигмоид на ReLU на последнем слое, но безрезультатно.
Это функция потерь и оптимизаторы, которые мы использовали:

Код: Выделить всё

optimizer = torch.optim.Adam(classifier.parameters(), lr=lr)
criterion = torch.nn.CrossEntropyLoss()

Вот пример ввода и вывода AE после обучения:



Подробнее здесь: https://stackoverflow.com/questions/667 ... from-color