Why is my Auto-Encoder not learning the distribution of FMNIST images?

desert_ranger: <p>I am using a simple autoencoder to learn images from the FashionMnist dataset. I have preprocessed the dataset by grayscaling and normalizing it. I did not make the network too deep, to prevent it from creating a direct mapping.</p>
<p>Here’s my PyTorch code -</p>
<pre><code class="lang-auto">import torch
import torchvision as tv
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from torch import nn
import os
from torchviz import make_dot
transforms = tv.transforms.Compose([tv.transforms.Grayscale(num_output_channels=1)])
trainset = tv.datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=transforms)
PATH = './ae.pth'
data = trainset.data.float()
data = data/255
# print(trainset.data.shape)
plt.imshow(trainset.data[0], cmap = 'gray')
plt.show()



class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.encode = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 30),
            nn.ReLU()
        )
        self.decode = nn.Sequential(
            nn.Linear(30, 512),
            nn.ReLU(),
            nn.Linear(512, 28*28),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.flatten(x)
        encoded = self.encode(x)
        decoded = self.decode(encoded)
        return decoded
if(os.path.exists(PATH)):
    print("Loading data on cpu")
    device = torch.device('cpu')
    model = NeuralNetwork()
    model.load_state_dict(torch.load(PATH, map_location=device))

else:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    data = data.to(device)
    print(f"Using device = {device}")
    model = NeuralNetwork().to(device)
    # print(model)

    lossFn  = nn.BCELoss()

    optimizer = torch.optim.SGD(model.parameters(), lr = 1e-3)

    for epoch in range(1000):
        print("Epoch = ", epoch)
        optimizer.zero_grad()
        outputs = model(data)
        loss = lossFn(outputs, data.reshape(-1, 784))
        loss.backward()
        optimizer.step()

    torch.save(model.state_dict(), PATH)
    data = data.to("cpu")
    model = model.to("cpu")

pred = model(data)
pred = pred.reshape(-1, 28, 28)
# print(pred.shape)
plt.imshow(pred.detach().numpy()[0], cmap = 'gray')
plt.show()
</code></pre>
<p>For testing, I am inputting the following image -<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/a/da0365e77647d8f872d5cd289915eccf74df908e.png" data-download-href="https://discuss.pytorch.org/uploads/default/da0365e77647d8f872d5cd289915eccf74df908e" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/a/da0365e77647d8f872d5cd289915eccf74df908e.png" alt="image" data-base62-sha1="v6DmRX73QeLtN1ooVjSdy5kZYN8" width="614" height="500" data-dominant-color="949494"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">646×526 7.77 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The following image gets outputted -<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/9/19ae0e25ae7193a1a08e041b4433a2a665282b49.png" data-download-href="https://discuss.pytorch.org/uploads/default/19ae0e25ae7193a1a08e041b4433a2a665282b49" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/9/19ae0e25ae7193a1a08e041b4433a2a665282b49.png" alt="image" data-base62-sha1="3FaQBpiBuY1stD1XgCYsY7fJshj" width="606" height="500" data-dominant-color="9E9E9E"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">633×522 10.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

omid_khalafbeigi: <p>Hi.<br>
When your model cant re-create input images, that means it can’t learn to extract features of image or decode them.<br>
I suggest you using layers like <a href="https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html" rel="noopener nofollow ugc">Convolution</a> and <a href="https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html" rel="noopener nofollow ugc">Deconvolution</a>.</p>