No Clue: KL divergence in VAE - IndexError: Dimension out of range

username: samim
cooked: <p>I am trying to implement a variational autoencoder, but calculating the Kullback Leibler divergence doesn’t work out the way I hoped.</p>
<p>Background: The input is a 1x800 tensor, which will be mapped to a 1x1200 tensor.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
from torch.utils.data import DataLoader, TensorDataset, IterableDataset
import os
from torch import optim
from model import Network

class Network(nn.Module):
    def __init__(self, input_dim, output_dim, latent_dim, layer_dim1, layer_dim2):
        """
        Parameter:
        input_dim (int): number of inputs
        output_dim (int): number of outputs
        latent_dim (int): number of latent neurons
        Layer_dim (int): number of neurons in hidden layers
        """
        super(Network, self).__init__()
        self.latent_dim = latent_dim

        self.enc1 = nn.Linear(input_dim, layer_dim1)
        self.enc2 = nn.Linear(layer_dim1, layer_dim2)

        self.latent = nn.Linear(layer_dim2, latent_dim*2)

        self.dec1 = nn.Linear(latent_dim, layer_dim2)
        self.dec2 = nn.Linear(layer_dim2, layer_dim1)

        self.out = nn.Linear(layer_dim1, output_dim)

    def encoder(self, x):
        z = F.elu(self.enc1(x))
        z = F.elu(self.enc2(x))
        z = self.latent(z)
        self.mu = z[0:self.latent_dim]
        self.log_sigma = z[self.latent_dim:]
        self.sigma = torch.exp(self.log_sigma)

        eps = torch.randn(x.size(0), self.latent_dim)
        z_sample = self.mu + self.sigma*eps

        self.kl_loss = kl_divergence(self.mu, self.log_sigma, dim=self.latent_dim)

        return z_sample

    def decoder(self, z):
        x = F.elu(self.dec1(z))
        x = F.elu(self.dec2(x))
        return self.out(x)

    def forward(self, batch):
        self.latent_rep = self.encoder(batch)
        dec_input = self.latent_rep
        return self.decoder(dec_input)


def kl_divergence(means, log_sigma, dim, target_sigma=0.1):
    """
    Computes Kullback–Leibler divergence for arrays of mean and log(sigma)
    """
    target_sigma = torch.Tensor([target_sigma])
    out = 1 / 2. * torch.mean(torch.mean(1 / target_sigma**2 * means**2 + torch.exp(2 * log_sigma) / target_sigma**2 - 2 * log_sigma + 2 * torch.log(target_sigma), dim=1) - dim)
    out = out
    return out

model = Network(800,1200,3,800,200)
SAVE_PATH = "trained/model.dat"
epochs = 5
learning_rate = 0.001
optimizer = optim.Adam(model.parameters(),lr=learning_rate, eps=1e-08)
hist_error = []
hist_loss = []
beta = 0.5



for epoch in range(epochs):
    epoch_error = []
    epoch_loss = []
    for i in x:
        optimizer.zero_grad()
        print(i)
        #i = torch.tensor(i)
        #i.unsqueeze_(0)
        pred = model.forward(i)
        loss = torch.mean(torch.sum((pred - y[i]) ** 2))
        loss.backward()
        optimizer.step()
        error = torch.mean(torch.sqrt((pred - y[i]) ** 2)).detach().numpy()
        epoch_error.append(error)
        epoch_loss.append(loss.data.detach().numpy())
    hist_error.append(np.mean(epoch_error))
    hist_loss.append(np.mean(epoch_loss))
    print("Epoch %d -- loss %f, RMS error %f " % (epoch+1, hist_loss[-1], hist_error[-1]))

</code></pre>
<p>The error message I now get is something I looked into last night, but could not resolve:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/samim/miniconda3/envs/deep/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 3343, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "&lt;ipython-input-2-54864ad18480&gt;", line 1, in &lt;module&gt;
    runfile('/home/samim/Documents/train.py', wdir='/home/samim/Documents')
  File "/home/samim/.local/share/JetBrains/PyCharm2020.3/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 197, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/home/samim/.local/share/JetBrains/PyCharm2020.3/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/home/samim/Documents/train.py", line 65, in &lt;module&gt;
    pred = model.forward(i)
  File "/home/samim/Documents/model.py", line 49, in forward
    self.latent_rep = self.encoder(batch)
  File "/home/samim/Documents/model.py", line 39, in encoder
    self.kl_loss = kl_divergence(self.mu, self.log_sigma, dim=self.latent_dim)
  File "/home/samim/Documents/model.py", line 59, in kl_divergence
    out = 1 / 2. * torch.mean(torch.mean(1 / target_sigma**2 * means**2 + torch.exp(2 * log_sigma) / target_sigma**2 - 2 * log_sigma + 2 * torch.log(target_sigma), dim=1) - dim)
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
</code></pre>

username: crcrpar
cooked: <p>Hi,</p>
<aside class="quote no-group" data-username="samim" data-post="1" data-topic="114554">
<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/samim/48/30150_2.png" class="avatar"> samim:</div>
<blockquote>
<p><code>out = 1 / 2. * torch.mean(torch.mean(1 / target_sigma**2 * means**2 + torch.exp(2 * log_sigma) / target_sigma**2 - 2 * log_sigma + 2 * torch.log(target_sigma), dim=1) - dim)</code></p>
</blockquote>
</aside>
<p>What’s the <code>ndim</code> of <code>1 / target_sigma**2 * means**2 + torch.exp(2 * log_sigma) / target_sigma**2 - 2 * log_sigma + 2 * torch.log(target_sigma)</code> which is the argument of innermost <code>torch.mean</code>?<br>
If <code>self.latent_dim</code> is greater than the <code>ndim</code> above, that error could happen.</p>

username: samim
cooked: <p>ndim is 3, as is self.latent_dim</p>
<p>when the inner torch.mean is executed with dim = 1, it outputs the error</p>
<pre><code class="lang-auto">IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
</code></pre>
<p>if it is 0, it’s all fine. my guess is that there is only 1 dimension (dim = 0) since the latent tensor is 1x3. Am I correct?</p>

username: patrickwilliams3
cooked: <aside class="quote no-group" data-username="samim" data-post="1" data-topic="114554">
<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/samim/48/30150_2.png" class="avatar"> samim:</div>
<blockquote>
<pre><code class="lang-auto">        self.mu = z[0:self.latent_dim]
        self.log_sigma = z[self.latent_dim:]
</code></pre>
</blockquote>
</aside>
<p>It looks like when you’re setting the self.mu and self.log_sigma you are sampling from the first dimension which is not the feature dimension but the batch size since the input tensor is a 1x800. I believe you need to sample from the second dimension here so try changing this to</p>
<pre><code class="lang-auto">        self.mu = z[:,0:self.latent_dim]
        self.log_sigma = z[:,self.latent_dim:]
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/8/d8ed5fcfc4f5a6ca80b6452a00b7413f52b553d7.png" data-download-href="https://discuss.pytorch.org/uploads/default/d8ed5fcfc4f5a6ca80b6452a00b7413f52b553d7" title="Screen Shot 2021-03-13 at 11.42.00 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/8/d8ed5fcfc4f5a6ca80b6452a00b7413f52b553d7_2_517x309.png" alt="Screen Shot 2021-03-13 at 11.42.00 AM" data-base62-sha1="uX1HKrUXViTIBWryBNgMaGujKvl" width="517" height="309" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/8/d8ed5fcfc4f5a6ca80b6452a00b7413f52b553d7_2_517x309.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/8/d8ed5fcfc4f5a6ca80b6452a00b7413f52b553d7_2_775x463.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/8/d8ed5fcfc4f5a6ca80b6452a00b7413f52b553d7_2_1034x618.png 2x" data-dominant-color="F5F3F3"><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">Screen Shot 2021-03-13 at 11.42.00 AM</span><span class="informations">1644×984 90.7 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: samim
cooked: <p>That was quite helpful! Thank you!</p>

username: muammar
cooked: <aside class="quote no-group" data-username="samim" data-post="1" data-topic="114554">
<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/samim/48/30150_2.png" class="avatar"> samim:</div>
<blockquote>
<p><code> kl_divergence</code></p>
</blockquote>
</aside>
<p>How did you derive this? and why did you not use the provided functional form of KL divergence? This is out of curiosity because I am looking to implement the VAE using <code>torch.nn.functional.kl_div</code>.</p>