How to implement validation loss in VAE training session

username: anu
cooked: <p>I am trying to calculate the validation loss along with the training set loss while training a simple VAE network. But, I am receiving a CUDA error as follows:</p>
<pre><code class="lang-auto">&lt;ipython-input-27-5638e9c724e6&gt; in unSupTrain(epoch)
     22             recon_batch, mu, sigma = model(data)
     23             # Get valloss value
---&gt; 24             vallossitem = elbo(recon_batch, data, mu, sigma).item()
     25             # Append loss to history
     26             hist_validation += vallossitem

RuntimeError: CUDA error: device-side assert triggered
</code></pre>
<p>It probably occurs because of the loss function setup but I somehow couldn’t find any solution for separating training and validation loss.</p>
<p>Here is the elbo function:</p>
<pre><code class="lang-auto">def elbo(recon_x, x, mu, sigma):
    '''Loss function.'''
    # Reshape the input
    x = x.view(-1, INP_SIZE)

    # Binary cross entropy
    RE = F.binary_cross_entropy(recon_x, x, reduction='sum')
    # KL divergence
    KL = F.kl_div(recon_x, x, reduction='sum', log_target=True)
    # Return the loss
    return RE - KL
</code></pre>
<p>and training function (certain parts of validation and training loss calculation):</p>
<pre><code class="lang-auto"># ---------- Validation -----------
    # Don't track gradients
    with torch.no_grad():
        # Trace
        print('Running Validation')
        for batch_idx, (data, _) in enumerate(valloader):
            # Convert to cuda if possible
            data = data.to(device=DEVICE)

            # Forward
            recon_batch, mu, sigma = model(data)
            # Get valloss value
            vallossitem = elbo(recon_batch, data, mu, sigma).item()
            # Append loss to history
            hist_validation += vallossitem
            # Round loss
            vallossitem = round(vallossitem / len(data), 3)

# -------- Training -------- #
    for batch_idx, (data, _) in enumerate(trainloader):
        # Convert to cuda if possible
        data = data.to(device=DEVICE)

        # Zero grad
        optimizer.zero_grad()
        # Forward
        recon_batch, mu, sigma = model(data)
        # Loss calculation
        trainloss = elbo(recon_batch, data, mu, sigma)
        # Backward
        trainloss.backward()
        # Get trainloss value
        trainlossitem = trainloss.item()
        # Append loss to history
        hist_training += trainlossitem
</code></pre>
<p>I would appreciate any suggestions. Thanks in advance!</p>

username: ptrblck
cooked: <p>Could you rerun the code via:</p>
<pre><code class="lang-auto">CUDA_LAUNCH_BLOCKING=1 python script.py args
</code></pre>
<p>and post the complete stack trace here?<br>
Alternatively, you could also run the script on the CPU, which might yield a better error message.</p>

username: muammar
cooked: <p>I doubted this <code>elbo</code> function. Usually, in the examples, we have the loss function implemented in the following way:</p>
<pre><code>def loss_function(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction="sum")

    # see Appendix B from VAE paper:
    #  Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # https://arxiv.org/abs/1312.6114
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    return BCE + KLD, BCE, KLD
</code></pre>
<p>I wanted to reproduce this using the <code>torch.nn.functional.kl_div</code> function, which is similar to what OP implemented in <code>elbo</code>:</p>
<pre><code>def loss_function(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction="sum")

    # see Appendix B from VAE paper:
    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # https://arxiv.org/abs/1312.6114
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    
    KLD = torch.nn.functional.kl_div(
        recon_x, 
        x.view(-1, 784), 
        reduction='batchmean', 
        log_target=False
    )

    return BCE + KLD, BCE, KLD
</code></pre>
<p>But, I wonder — Is this correct? I ask that because the KLD should measure the divergence between the mean of the posterior (latent distribution) with respect to the mean of the prior (normal distribution). One can see that this does not work because reconstructions are sharp, but the generative aspect of the VAE is completely lost compared to the loss when you use</p>
<p><code>KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())</code>.</p>
<p>I think that <code>kl_div</code> should be receiving <code>mu</code> that comes from the latent distribution, and for the <code>target</code>, it should be receiving the <code>eps</code> from the <code>reparameterize</code> function?</p>
<pre><code>def reparameterize(self, mu, logvar):
    
    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)

    return mu + eps * std, eps
</code></pre>
<p>PS. I probably should have opened a new post for this, but it made sense here in my mind.</p>

username: WhyNotML
cooked: <p>Should the output be given to the loss_function?<br>
<code>loss_function(output)</code>? Did I get that right? The output I have received is a tuple that has three tensor elements with [1,3,32,32], 2x [1,512]</p>
<p>I couldn’t pass the output directly, I think because it is a tuple.</p>