<p>Thanks for providing the code and a link to a Colab notebook! +1!
Also, your code is well-written and easy to read. Unless I missed something, I think there are two problems with your code:</p>
<ol>
<li>The data normalization</li>
<li>The implementation of the VAE loss.</li>
</ol>
<p><strong>About 1.</strong>, your <code>CIFAR10DataModule</code> class normalizes the RGB channels of the CIFAR10 images using <code>mean = 0.5</code> and <code>std = 0.5</code>. Since the pixel values are initially in [0,1] range, the normalized images have pixel values in the [-1,1] range. However, your <code>Decoder</code> class applies a <code>nn.Sigmoid()</code> activation to the reconstructed images. Therefore, your reconstructed images have pixel values in the [0,1] range. I suggest to remove this mean-std normalization so that both the &quot;true&quot; images and the reconstructed images have their pixel values in the [0,1] range.</p>
<p><strong>About 2.</strong>: since you're dealing with RGB images the MSE loss makes sense. The idea behind the MSE loss is the &quot;Gaussian decoder&quot;. This decoder assumes the pixel values of a &quot;true image&quot; is generated by independent Gaussian distributions whose mean is the pixel values of the reconstructed image (i.e. the output of the decoder) and with a given variance. Your implementation of the reconstruction loss (namely <code>r_loss = F.mse_loss(predictions, targets)</code>) is equivalent to a fixed variance. Using ideas from <a href="https://arxiv.org/abs/2006.13202" rel="nofollow noreferrer">this paper</a>, we can do better and obtain an analytic expression for the &quot;optimal value&quot; of this variance parameter. Finally, the reconstruction loss should be summed over all pixels (<code>reduction = 'sum'</code>). To understand why, have a look at analytic expression of the reconstruction loss (see, for instance, <a href="http://adamlineberry.ai/vae-series/vae-code-experiments" rel="nofollow noreferrer">this blog post</a> which considers the BCE loss).</p>
<p>Here is what the refactored <code>LitVAE</code> class looks like:</p>
<pre class="lang-py prettyprint-override"><code>class LitVAE(pl.LightningModule):
    def __init__(self,
                 learning_rate: float = 0.0005,
                 **kwargs) -&gt; None:
        &quot;&quot;&quot;
        Parameters
        ----------
        - `learning_rate: float`:
            learning rate for the optimizer
        - `**kwargs`:
            arguments to pass to the variational autoencoder constructor
        &quot;&quot;&quot;
        super(LitVAE, self).__init__()
        
        self.learning_rate = learning_rate 

        self.vae = VariationalAutoEncoder(**kwargs)

    def forward(self, x) -&gt; _tensor_size_3_t: 
        return self.vae(x)

    def training_step(self, batch, batch_idx):
        r_loss, kl_loss, sigma_opt = self.shared_step(batch)
        loss = r_loss + kl_loss
        
        self.log(&quot;train_loss_step&quot;, loss)
        return {&quot;loss&quot;: loss, 'log':{&quot;r_loss&quot;: r_loss / len(batch[0]), &quot;kl_loss&quot;: kl_loss / len(batch[0]), 'sigma_opt': sigma_opt}}

    def training_epoch_end(self, outputs) -&gt; None:
        # add computation graph
        if(self.current_epoch == 0):
            sample_input = torch.randn((1, 3, 32, 32))
            sample_model = LitVAE(**MODEL_PARAMS)
            
            self.logger.experiment.add_graph(sample_model, sample_input)
            
        epoch_loss = self.average_metric(outputs, &quot;loss&quot;)
        self.logger.experiment.add_scalar(&quot;train_loss_epoch&quot;, epoch_loss, self.current_epoch)

    def validation_step(self, batch, batch_idx):
        r_loss, kl_loss, _ = self.shared_step(batch)
        loss = r_loss + kl_loss

        self.log(&quot;valid_loss_step&quot;, loss)

        return {&quot;loss&quot;: loss}

    def validation_epoch_end(self, outputs) -&gt; None:
        epoch_loss = self.average_metric(outputs, &quot;loss&quot;)
        self.logger.experiment.add_scalar(&quot;valid_loss_epoch&quot;, epoch_loss, self.current_epoch)

    def test_step(self, batch, batch_idx):
        r_loss, kl_loss, _ = self.shared_step(batch)
        loss = r_loss + kl_loss
        
        self.log(&quot;test_loss_step&quot;, loss)
        return {&quot;loss&quot;: loss}

    def test_epoch_end(self, outputs) -&gt; None:
        epoch_loss = self.average_metric(outputs, &quot;loss&quot;)
        self.logger.experiment.add_scalar(&quot;test_loss_epoch&quot;, epoch_loss, self.current_epoch)

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.learning_rate)
        
    def shared_step(self, batch) -&gt; torch.TensorType: 
        # images are both samples and targets thus original 
        # labels from the dataset are not required
        true_images, _ = batch

        # perform a forward pass through the VAE 
        # mean and log_variance are used to calculate the KL Divergence loss 
        # decoder_output represents the generated images 
        mean, log_variance, generated_images = self(true_images)

        r_loss, kl_loss, sigma_opt = self.calculate_loss(mean, log_variance, generated_images, true_images)
        return r_loss, kl_loss, sigma_opt

    def calculate_loss(self, mean, log_variance, predictions, targets):
        mse = F.mse_loss(predictions, targets, reduction='mean')
        log_sigma_opt = 0.5 * mse.log()
        r_loss = 0.5 * torch.pow((targets - predictions) / log_sigma_opt.exp(), 2) + log_sigma_opt
        r_loss = r_loss.sum()
        kl_loss = self._compute_kl_loss(mean, log_variance)
        return r_loss, kl_loss, log_sigma_opt.exp()

    def _compute_kl_loss(self, mean, log_variance): 
        return -0.5 * torch.sum(1 + log_variance - mean.pow(2) - log_variance.exp())

    def average_metric(self, metrics, metric_name):
        avg_metric = torch.stack([x[metric_name] for x in metrics]).mean()
        return avg_metric
</code></pre>
<p>After 10 epochs, that's what the reconstructed images look like:</p>
<p><a href="https://i.stack.imgur.com/c0K3W.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/c0K3W.png" alt="enter image description here" /></a></p>
What could cause a VAE(Variational AutoEncoder) to output random noise even after training?
<p>I have trained a VAE on CIFAR10 data-set. However, when I try to generate images from the VAE all I get is a bunch of gray noise back. The implementation of this VAE follows the implementation from the book <a href="https://www.oreilly.com/library/view/generative-deep-learning/9781492041931/" rel="nofollow noreferrer">Generative Deep Learning</a>, but instead of TensorFlow the code uses PyTorch.</p>
<p>The notebook containing the training as well as the generation can be found <a href="https://colab.research.google.com/drive/1t0EmOo5UjmeRScFtQzpcDEB8u2vYW7Yk?usp=sharing" rel="nofollow noreferrer">here</a>, while the actual implementation of the VAE can be found <a href="https://github.com/Kinyugo/generative-deep-learning/blob/main/variational_autoencoder/variational_autoencoder.py" rel="nofollow noreferrer">here</a>.</p>
<p>I have tried:</p>
<ol>
<li>Disabling dropouts.</li>
<li>Increasing the dimension of the latent space.</li>
</ol>
<p>None of the methods show any improvement at all.</p>
<p>I have verified that:</p>
<ol>
<li>The input size matches the output size</li>
<li>Back-propagation runs successfully as the loss decreases during training.</li>
</ol>