Convolutional VAE not training

username: tirthasheshpatel
cooked: <p>I have been trying to implement convolutional VAE in PyTorch for a while now and am somehow not able to correctly train my network. Here’s the encoder, decoder, and training loop. I am training the model on MNIST dataset.</p>
<p>Encoder output: Two tensors (loc, logvar) of shape <code>[batch_size, latent_dims]</code></p>
<p>Decoder output: Image of shape <code>[batch_size, 1, 28, 28]</code></p>
<p>Problem: Loss remains almost constant.</p>
<details>
<summary>Encoder</summary>
<pre><code class="lang-python">class LocLogvar(nn.Module):
    def __init__(self, in_features, latent_dims):
        super(LocLogvar, self).__init__()
        self.loc = nn.Linear(in_features, latent_dims)
        self.logvar = nn.Linear(in_features, latent_dims)

    def forward(self, inputs):
        loc = self.loc(inputs)
        logvar = self.logvar(inputs)
        return loc, logvar

encoder = nn.Sequential(OrderedDict([
        ('e_conv_layer_1', nn.Conv2d(1, 16, 5, 1)),                # 16 x 24 x 24
        ('e_relu_layer_1', IPLReLU()),
        ('e_batch_norm_1', nn.BatchNorm2d(16)),
        ('e_conv_layer_2', nn.Conv2d(16, 32, 5, 1)),               # 32 x 20 x 20
        ('e_relu_layer_2', IPLReLU()),
        ('e_batch_norm_2', nn.BatchNorm2d(32)),
        ('e_conv_layer_3', nn.Conv2d(32, 32, 11, 1)),              # 32 x 10 x 10
        ('e_relu_layer_3', IPLReLU()),
        ('e_batch_norm_3', nn.BatchNorm2d(32)),
        ('e_conv_layer_4', nn.Conv2d(32, 64, 5, 1)),               # 64 x 6 x 6
        ('e_relu_layer_4', IPLReLU()),
        ('e_batch_norm_4', nn.BatchNorm2d(64)),
        ('e_dropout_layer_1', nn.Dropout2d(p=0.75)),
        ('e_conv_layer_5', nn.Conv2d(64, 128, 5, 1)),              # 128 x 2 x 2
        ('e_relu_layer_5', IPLReLU()),
        ('e_batch_norm_5', nn.BatchNorm2d(128)),
        ('e_dropout_layer_1', nn.Dropout2d(p=0.85)),
        ('e_flatten_layer', nn.Flatten()),
        ('out_layer', LocLogvar(128*2*2, latent_dims))
]))
</code></pre>
</details>
<details>
<summary>Decoder</summary>
<pre><code class="lang-python">class Reshape(nn.Module):
    def __init__(self, *shape):
        super(Reshape, self).__init__()
        self.shape = shape

    def forward(self, X):
        return X.view(-1, *self.shape)

decoder = nn.Sequential(OrderedDict([
        ('inv_linear_layer_1', nn.Linear(latent_dims, 128*2*2)),   # 128 * 2 * 2
        ('inv_relu_layer_5', IPLReLU()),
        ('inv_flatten_layer', Reshape(128, 2, 2)),                 # 128 x 2 x 2
        ('inv_conv_layer_5', nn.ConvTranspose2d(128, 64, 5, 1)),   # 64 x 6 x 6
        ('inv_batch_norm_5', nn.BatchNorm2d(64)),
        ('inv_relu_layer_4', IPLReLU()),
        ('inv_conv_layer_4', nn.ConvTranspose2d(64, 32, 5, 1)),    # 32 x 10 x 10
        ('inv_batch_norm_4', nn.BatchNorm2d(32)),
        ('inv_relu_layer_3', IPLReLU()),
        ('inv_conv_layer_3', nn.ConvTranspose2d(32, 32, 11, 1)),   # 32 x 20 x 20
        ('inv_batch_norm_3', nn.BatchNorm2d(32)),
        ('inv_relu_layer_2', IPLReLU()),
        ('inv_conv_layer_2', nn.ConvTranspose2d(32, 16, 5, 1)),    # 16 x 24 x 24
        ('inv_batch_norm_2', nn.BatchNorm2d(16)),
        ('inv_relu_layer_1', IPLReLU()),
        ('inv_conv_layer_1', nn.ConvTranspose2d(16, 1, 5, 1)),     # 1 x 28 x 28
        ('inv_batch_norm_1', nn.BatchNorm2d(1)),
        ('out_layer', nn.Sigmoid())
]))
</code></pre>
</details>
<details>
<summary>Loss function</summary>
<pre><code class="lang-python">def loss_fn(loc, logvar, reconstructed, img):
    reconstructed = reconstructed.view(-1, 784)
    img = img.view(-1, 784)
    recon_loss = -torch.sum(
        img*torch.log(reconstructed+1e-18) + (1-img)*torch.log(1-reconstructed+1e-18)
    )
    kl_loss = 0.5 * torch.sum(
        -logvar - 1 + logvar.exp() + loc**2
    )
    return recon_loss + kl_loss
</code></pre>
</details>
<details>
<summary>Training Loop</summary>
<pre><code class="lang-python">encoder = encoder.to(torch.device('cuda'))
decoder = decoder.to(torch.device('cuda'))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for i in range(epochs):
    for idx in range(x_train.shape[0] // batch_size):
        # Get the batch to train
        x_batch = x_train[idx*batch_size:(idx+1)*batch_size, ...]

        # Forward pass through encoder
        loc, logvar = encoder(x_batch)

        # Reparameterize
        epsilon = torch.randn_like(loc)
        z = loc + torch.exp(logvar * 0.5) * epsilon

        # Forward pass through decoder
        reconstructed = decoder(z)
        loss = loss_fn(loc, logvar, reconstructed, x_batch)

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

        # Print the results
        stdout.write(f'\r epoch : {i}\t'
                     f'step : {min((idx+1)*batch_size, x_train.shape[0])}/{x_train.shape[0]}\t'
                     f'loss : {loss.item():.3f}\t')
    print()
</code></pre>
</details>
<details>
<summary>Output</summary>
<pre><code class="lang-python"> epoch : 0	step : 60000/60000	loss : 634711.188	
 epoch : 1	step : 60000/60000	loss : 635014.250	
 epoch : 2	step : 60000/60000	loss : 634935.625	
 epoch : 3	step : 60000/60000	loss : 635042.812	
 epoch : 4	step : 60000/60000	loss : 634130.562	
 epoch : 5	step : 60000/60000	loss : 634780.438	
 epoch : 6	step : 60000/60000	loss : 634427.250	
 epoch : 7	step : 60000/60000	loss : 634962.000	
 epoch : 8	step : 60000/60000	loss : 634118.750	
 epoch : 9	step : 60000/60000	loss : 636151.125
</code></pre>
</details>
<p>I suspect some gradients are getting detached somewhere but I am not sure what is causing this behaviour. It would be great if someone can suggest my mistake in the model <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: ptrblck
cooked: <p>I can’t spot any obvious error or a line of code, which might detach the tensors from the computation graph.<br>
Just to make sure, you are indeed not detaching, you could print the gradients of all parameters after the backward pass:</p>
<pre><code class="lang-python">for name, param in encoder.named_parameters():
    print(name, param.grad)
# same for decoder
</code></pre>
<p>If you see some <code>None</code> outputs, this would mean that these parameters won’t get gradients and thus won’t be updated.</p>

username: tirthasheshpatel
cooked: <p>Hello <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>First of all, thanks for your reply! I checked if any gradients were none, infinity or nan. Turns out that’s not the case still the model doesn’t train. But when I create a module with encoder and decoder in the module as submodules, it magically (for me) starts training as expected and produces really nice samples!! But I don’t get it. Shouldn’t the model train even if I have encoder and decoder as seperate modules and not shoved in <code>nn.Module</code> class? I may be wrong there as I have just started using torch.</p>
<p>My modified code:</p>
<details>
<summary>Modified model</summary>
<pre><code class="lang-python">class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()

        self.encoder = encoder = nn.Sequential(OrderedDict([
        ('e_conv_layer_1', nn.Conv2d(1, 16, 5, 1)),                # 16 x 24 x 24
        ('e_relu_layer_1', nn.LeakyReLU(inplace=True)),
        ('e_batch_norm_1', nn.BatchNorm2d(16)),
        ('e_conv_layer_2', nn.Conv2d(16, 32, 5, 1)),               # 32 x 20 x 20
        ('e_relu_layer_2', nn.LeakyReLU(inplace=True)),
        ('e_batch_norm_2', nn.BatchNorm2d(32)),
        ('e_conv_layer_3', nn.Conv2d(32, 32, 11, 1)),              # 32 x 10 x 10
        ('e_relu_layer_3', nn.LeakyReLU(inplace=True)),
        ('e_batch_norm_3', nn.BatchNorm2d(32)),
        ('e_conv_layer_4', nn.Conv2d(32, 64, 5, 1)),               # 64 x 6 x 6
        ('e_relu_layer_4', nn.LeakyReLU(inplace=True)),
        ('e_batch_norm_4', nn.BatchNorm2d(64)),
        ('e_dropout_layer_1', nn.Dropout2d(p=0.75)),
        ('e_conv_layer_5', nn.Conv2d(64, 128, 5, 1)),              # 128 x 2 x 2
        ('e_relu_layer_5', nn.LeakyReLU(inplace=True)),
        ('e_batch_norm_5', nn.BatchNorm2d(128)),
        ('e_dropout_layer_1', nn.Dropout2d(p=0.85)),
        ('e_flatten_layer', nn.Flatten()),
        ('e_out_layer', LocLogvar(128*2*2, latent_dims))
        ]))

        self.decoder = decoder = nn.Sequential(OrderedDict([
        ('inv_linear_layer_1', nn.Linear(latent_dims, 128*2*2)),   # 128 * 2 * 2
        ('inv_relu_layer_5', nn.LeakyReLU(inplace=True)),
        ('inv_flatten_layer', Reshape(128, 2, 2)),                 # 128 x 2 x 2
        ('inv_conv_layer_5', nn.ConvTranspose2d(128, 64, 5, 1)),   # 64 x 6 x 6
        ('inv_batch_norm_5', nn.BatchNorm2d(64)),
        ('inv_relu_layer_4', nn.LeakyReLU(inplace=True)),
        ('inv_conv_layer_4', nn.ConvTranspose2d(64, 32, 5, 1)),    # 32 x 10 x 10
        ('inv_batch_norm_4', nn.BatchNorm2d(32)),
        ('inv_relu_layer_3', nn.LeakyReLU(inplace=True)),
        ('inv_conv_layer_3', nn.ConvTranspose2d(32, 32, 11, 1)),   # 32 x 20 x 20
        ('inv_batch_norm_3', nn.BatchNorm2d(32)),
        ('inv_relu_layer_2', nn.LeakyReLU(inplace=True)),
        ('inv_conv_layer_2', nn.ConvTranspose2d(32, 16, 5, 1)),    # 16 x 24 x 24
        ('inv_batch_norm_2', nn.BatchNorm2d(16)),
        ('inv_relu_layer_1', nn.LeakyReLU(inplace=True)),
        ('inv_conv_layer_1', nn.ConvTranspose2d(16, 1, 5, 1)),     # 1 x 28 x 28
        ('inv_batch_norm_1', nn.BatchNorm2d(1)),
        ('inv_out_layer', nn.Sigmoid())
        ]))

    def encode(self, x):
        return self.encoder(x)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std

    def decode(self, z):
        return self.decoder(z)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

</code></pre>
</details>

username: ptrblck
cooked: <aside class="quote no-group" data-username="tirthasheshpatel" data-post="3" data-topic="83187">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tirthasheshpatel/48/24845_2.png" class="avatar"> tirthasheshpatel:</div>
<blockquote>
<p>Shouldn’t the model train even if I have encoder and decoder as seperate modules and not shoved in <code>nn.Module</code> class? I may be wrong there as I have just started using torch.</p>
</blockquote>
</aside>
<p>You are right and it shouldn’t make a difference, as the computation graph will be the same.</p>
<p>It might be a typo or copy-paste issue, but your <code>optimizer</code> in the first approach takes the parameters from <code>model</code>, which is undefined. It should get the parameters of <code>encoder</code> and <code>decoder</code> so that it can update these parameters using their gradients.<br>
Could you check this particular line of code?</p>
<p>If this is indeed a typo, could you try to rerun both approaches for a couple of epochs using different seeds and check how reproducible the effect is, that the second approach trains better than the first?</p>

username: tirthasheshpatel
cooked: <p>Oh!!! That was a silly silly mistake! How did I not notice that!! So sorry to take your time. Thanks for the help! Cheers!</p>