Unable to train VAE Decoder

username: CaglarDeniz
cooked: <p>Hello Everyone,<br>
I’m hesitant posting this here as I’m not sure if this forum is open to debugging questions. Nevertheless, I’m willing to take my chances.</p>
<p>I have followed <a href="https://medium.com/dataseries/variational-autoencoder-with-pytorch-2d359cbf027b" rel="noopener nofollow ugc">this</a> tutorial right here to build my own VAE, which I train on my custom data set with the same type of images as MNIST which is used on the tutorial(28x28, black and white).</p>
<p>However, I must have done something wrong while implementing it, since my model won’t learn anything. It just outputs a specific pattern for each input image ,as you can see below, even though the input images have huge differences.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/d/fd5bf5244cc7a99a166b70b580bbc8e2e1fd0a09.png" data-download-href="https://discuss.pytorch.org/uploads/default/fd5bf5244cc7a99a166b70b580bbc8e2e1fd0a09" title="Output after 100 epochs"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/d/fd5bf5244cc7a99a166b70b580bbc8e2e1fd0a09_2_384x375.png" alt="Output after 100 epochs" data-base62-sha1="A9jOJjM1rIG5Mkes2Ag6qIWIeeJ" width="384" height="375" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/d/fd5bf5244cc7a99a166b70b580bbc8e2e1fd0a09_2_384x375.png, https://discuss.pytorch.org/uploads/default/original/3X/f/d/fd5bf5244cc7a99a166b70b580bbc8e2e1fd0a09.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/d/fd5bf5244cc7a99a166b70b580bbc8e2e1fd0a09.png 2x" data-dominant-color="E9E9E9"><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">Output after 100 epochs</span><span class="informations">554×539 20.2 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>This problem persists and the model insists even more so on this random pattern and just disregards the input.</p>
<p>I’ve included my code below. I can say that I have verified a couple things about my implementation after seeing people <a href="https://discuss.pytorch.org/t/convolutional-vae-not-training/83187">having similar issues</a>. I’m pretty sure all my gradients are updating. I also checked to see where this pattern is coming from and it seems like the encoder part is doing its job really well i.e. it outputs different encodings for different images. For some reason, the decoder just interprets these different encodings to very similar output images.</p>
<p>Any kind of suggestion would be much appreciated, and I am ready to supply any kind of debugging information you might request. Thank you in advance<br>
<strong>Encoder</strong></p>
<pre><code class="lang-auto"> class VariationalEncoder(nn.Module):
    def __init__(self, latent_dimensions):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(1, 8, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(8, 16, 3, stride=2, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.Conv2d(16,32,3, stride=2, padding=0),
            nn.ReLU(),
            nn.Flatten(start_dim=1),
            nn.Linear(3 * 3 * 32, 128),
            nn.ReLU(),
        )

        self.linear1 = nn.Linear(128, latent_dimensions)
        self.linear2 = nn.Linear(128, latent_dimensions)

        self.N = torch.distributions.Normal(0, 1)
        self.KLDivergence = 0

    def forward(self, x):
        out = self.model(x)

        mean = self.linear1(out)
        stddev = torch.exp(self.linear2(out))

        #  mu =  self.linear2(x)
        # sigma = torch.exp(self.linear3(x))
        # z = mu + sigma*self.N.sample(mu.shape)
        # self.kl = (sigma**2 + mu**2 - torch.log(sigma) - 1/2).sum()
        
        z = mean + stddev * self.N.sample(stddev.shape)

        self.KLDivergence =  (stddev**2 + mean**2 - torch.log(stddev) - 1/2).sum()
        return z

</code></pre>
<p><strong>Decoder</strong></p>
<pre><code class="lang-auto">class Decoder(nn.Module):
    def __init__(self, latent_dimensions):
        super().__init__()

        self.decoder_lin = nn.Sequential(
            nn.Linear(latent_dimensions, 128),
            nn.ReLU(True),
            nn.Linear(128, 3 * 3 * 32),
            nn.ReLU(True),
        )

        self.unflatten = nn.Unflatten(dim=1, unflattened_size=(32, 3, 3))

        self.decoder_conv = nn.Sequential(
            nn.ConvTranspose2d(32, 16, 3, stride=2, output_padding=0),
            nn.BatchNorm2d(16),
            nn.ReLU(True),
            nn.ConvTranspose2d(16, 8, 3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(8),
            nn.ReLU(True),
            nn.ConvTranspose2d(8, 1, 3, stride=2, padding=1, output_padding=1),
        )

    def forward(self, x):
        x = self.decoder_lin(x)
        x = self.unflatten(x)
        x = self.decoder_conv(x)
        x = torch.sigmoid(x)

        return x

</code></pre>
<p><strong>Parent Class</strong></p>
<pre><code class="lang-auto">class VariationalAutoEncoder(nn.Module):
    def __init__(self, latent_dimensions):
        super().__init__()
        self.encoder = VariationalEncoder(latent_dimensions)
        self.decoder = Decoder(latent_dimensions)

        # set default floating point data type to float64

        torch.set_default_dtype(torch.float64)

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

</code></pre>
<p><strong>Training Loop</strong></p>
<pre><code class="lang-auto">def train_epoch(vae, dataloader, optimizer):
    # Set train mode

    vae.train()
    vae.encoder.train()
    vae.decoder.train()
    train_loss = 0.0

    # set_trace()

    for x in tqdm(dataloader):
        # Get model output

        optimizer.zero_grad()
        # x = torch.from_numpy(x)
        x_hat = vae(x)

        # Evaluate loss
        loss = torch.sum((x - x_hat) ** 2) + vae.encoder.KLDivergence

        # Backward pass
        loss.backward()
        optimizer.step()

        # Print batch loss
        # print(f"\t partial training loss (single batch): {loss.item()} ")
        train_loss += loss.item()

    return train_loss / len(dataloader.dataset)
</code></pre>
<p>Thank you for your time!</p>

username: ptrblck
cooked: <p>A good first debugging step would be to try to overfit a small subset of the dataset (e.g. use only 10 samples) and check if your model is able to train on this dataset by playing around with some hyperparameters.<br>
If that’s not possible your training code might have other bugs which break the training.</p>

username: CaglarDeniz
cooked: <p>Thank you for your quick response! I tried what you recommended and trained on only 10 samples and ran 1000 epochs to make sure that I overfitted. Here’s what the model outputted for those 10 samples.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/9/9956627b638a618a14832622db6c7160eace6f68.png" data-download-href="https://discuss.pytorch.org/uploads/default/9956627b638a618a14832622db6c7160eace6f68" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/9/9956627b638a618a14832622db6c7160eace6f68.png" alt="image" data-base62-sha1="lSu88QKbSSyHjsJNUGc5Ryj23jy" width="557" height="500" data-dominant-color="E6E6E6"><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">569×510 15.8 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>

username: CaglarDeniz
cooked: <p>I just realized I didn’t include some of the information that you mentioned so I just wanted to add the remaining of my training code.</p>
<pre><code class="lang-auto">def train():
    # Set rand number seed for reproducible results

    torch.manual_seed(0)
    dataset = PopulationOneLineDataset("../pop_data")

    transform = transforms.Compose([transforms.ToTensor()])

    dataset.transform = transform

    m = len(dataset)

    # train_data, test_dataset = random_split(dataset, [m - int(m * 0.2), int(m * 0.2)])
    batch_size = 256

    train_loader = DataLoader(dataset, batch_size=batch_size,)
    # eval_loader = DataLoader(test_dataset, batch_size=batch_size,shuffle=True)

    # torch.manual_seed(0)
    d = 3
    vae = VariationalAutoEncoder(latent_dimensions=d)

    learning_rate = 1e-4

    optim = torch.optim.Adam(vae.parameters(), lr=learning_rate)

    training_loss = []

    num_epochs = 1000

    for epoch in tqdm(range(num_epochs)):
        train_loss = train_epoch(vae, train_loader, optim)
        training_loss.append(train_loss)
        # eval_loss = test_epoch(vae, eval_loader)
        # print('\n EPOCH {}/{} \t train loss {:.3f} \t val loss {:.3f}'.format(epoch + 1, num_epochs,train_loss,eval_loss))

    plot_ae_outputs(vae.encoder, vae.decoder, dataset,n=10)
    torch.save(vae, "vae.model")

    plt.plot(range(len(training_loss)),training_loss)
    plt.show()

</code></pre>

username: ptrblck
cooked: <p>Using your code I can properly overfit a single sample of <code>MNIST</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">device = "cpu"
model = VariationalAutoEncoder(64).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)


dataset = datasets.MNIST(root="~/python/data", download=False, transform=transforms.ToTensor())
x, y = dataset[0]
x = x.to(device)
x = x.unsqueeze(0)

for epoch in range(1000):
    optimizer.zero_grad()
    output = model(x)
    loss = torch.sum((x - output) ** 2) + model.encoder.KLDivergence
    loss.backward()
    optimizer.step()
    
    if epoch % 10 == 0:
        print(f"epoch: {epoch}, loss: {loss.item():.3}")
        plt.imshow(output[0, 0].detach().cpu().numpy())
</code></pre>
<p>Resulting in:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/2/6294d3b82cdbfec6dda3813112c353ee219d0564.png" alt="image" data-base62-sha1="e45EsniYffMBdytUnECgXdeECxe" width="251" height="248"><br>
to<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/e/3/e3eed2f9e43046aa8976651db784f509945c41c3.png" alt="image" data-base62-sha1="wwo3YtQrZFrJFVWrnIGvWwnLLFx" width="251" height="248"></p>
<p>The <code>KLDivergence</code> loss part seems to be stuck but I don’t know if it could be related to your issue in not being able to overfit a simple test case.</p>

username: CaglarDeniz
cooked: <p>I see, this makes me feel like something might be wrong with my dataset. I will look into it and report back here if I find something. Thank you so much for your time and efforts!</p>

username: hrbigelow
cooked: <p>Hi Deniz,</p>
<p>I’m not sure if this could be your issue, but a phenomenon called ‘posterior collapse’ can happen with the original VAE.  The ELBO objective amounts to - D[q(x,z) || p(x,z)] (see eq 2 of <a href="https://arxiv.org/abs/1706.02262" rel="noopener nofollow ugc">InfoVAE paper</a>) in the section “Equivalent Forms of the ELBO objective”.  What then happens is that during training, the VAE enters a very bad local minimum and cannot recover.  The local minimum is characterized by the decoder ignoring z (gradients w.r.t. z go to zero) and simply tries to fit the distribution of the dataset unconditionally.</p>
<p>This phenomenon is not well understood (by me!) but it surely depends on the structure of the dataset.  So, it is entirely possible that it would happen with your dataset but not with MNIST.</p>
<p>But did you get it to work on the MNIST dataset, just to make sure you implemented everything correctly according to the tutorial?</p>

username: CaglarDeniz
cooked: <p>Hi Henry,<br>
Thanks for taking the time to answer my question. I will look into this posterior collapse phenomenon, but I haven’t really worked on that dataset for a while. I was able to get it working on MNIST with the exact same code (only the data loading part of the code was different).</p>
<p>Best,<br>
Deniz</p>