自动编码器从彩色图像输出黑白图像

2024-09-30 08:27:37 发布

您现在位置:Python中文网/ 问答频道 /正文

我们正在培训CIFAR10图像的AE。我们使用了以下体系结构:

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替换Sigmoid,但没有成功

以下是我们使用的损失函数和优化器:

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

以下是训练后AE的输入和输出示例:

enter image description here


Tags: in图像selfencoderforinitdefnn

热门问题