No sample variety on MNIST for pl_bolts.models.autoencoders.VAE

4158ndfkvBHJ1: <p>I hope this is OK to post here, as it’s not directly related to PyTorch. I’ll delete it if it doesn’t fit.<br>
I am trying to train a variational autoencoder on small (64x64) grayscale patches. Ultimately, the downstream task is classification, but as most of my data is not labelled, semi-supervised learning seem like an interesting solution.<br>
I am using PyTorch lightning, but training a VAE on my images lead to absolutely 0 sample variety. I have the following code :</p>
<pre><code class="lang-auto">import os 
from torch.nn import Conv2d
from pytorch_lightning import Trainer
from pl_bolts.models.autoencoders import VAE 
from pl_bolts.datamodules import MNISTDataModule
from pl_bolts.callbacks.variational import LatentDimInterpolator as LDI

# interpolator callback  

interpolator = LDI(interpolate_epoch_interval = 1)

# training 

dataset = MNISTDataModule(num_workers = 8, batch_size = 128)
model = VAE(input_height = 28, enc_type = "resnet18", kl_coeff=1)

# changing first and last convolution to deal with 28x28 grayscale images 

model.encoder.conv1 = Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)
model.decoder.conv1 = Conv2d(64*model.decoder.expansion, 1, kernel_size=3, stride=1, padding=3, bias=False)
print(model)

# fit model 

trainer = Trainer(gpus=1,  max_epochs=50, callbacks = [interpolator], default_root_dir = os.getcwd())
trainer.fit(model, dataset)
trainer.save_checkpoint(os.getcwd())
</code></pre>
<p>I only modified the first and last convolution of the resnet decoder and encoder, because the images only have 1 channel. There are some small modifications in the callback code (I added a save directory to see the images as I’m not using tensorboard, and I changed a <code>self.module</code> call into <code>self.module.decoder()</code> as it was not running before).</p>
<p>This code gives me the following result (50 epochs, lr = 0.0001, Adam optimizer) :<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/9/a/9a036b0a961b35076cbd80f11bd69f42936cac2a.png" alt="3" data-base62-sha1="lYsQVQNsb2OeUhLF01xTLgTMgQy" width="302" height="392"></p>
<p>On my dataset (grayscale chromosome patches) I get similar results :<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/1/81406fc6a2d0607d3679118a121a183826b9a0e3.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/81406fc6a2d0607d3679118a121a183826b9a0e3" title="49"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81406fc6a2d0607d3679118a121a183826b9a0e3_2_384x500.jpeg" alt="49" data-base62-sha1="irpEcUo13HLQEb4yRazv0hf6tSX" width="384" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81406fc6a2d0607d3679118a121a183826b9a0e3_2_384x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81406fc6a2d0607d3679118a121a183826b9a0e3_2_576x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/1/81406fc6a2d0607d3679118a121a183826b9a0e3.jpeg 2x" data-dominant-color="202020"><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">49</span><span class="informations">662×860 79.3 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 post (<a href="https://www.quora.com/How-do-you-fix-a-Variational-Autoencoder-VAE-that-suffers-from-mode-collapse" rel="noopener nofollow ugc">https://www.quora.com/How-do-you-fix-a-Variational-Autoencoder-VAE-that-suffers-from-mode-collapse</a>) suggests that loss balancing could be an issue, but an hyperparameter search for kl_coeff didn’t turn anything up.</p>
<p>However, a fully connected VAE gives much better results, both on MNIST and on my chromosome dataset :</p>
<pre><code class="lang-auto">class encoder(nn.Module):

    def __init__(self, img_height, h_dim1, h_dim2, **kwargs) -&gt; None:

        super(encoder, self).__init__()

        self.fc1 = nn.Linear(img_height*img_height, h_dim1)
        self.fc2 = nn.Linear(h_dim1, h_dim2)

    def forward(self, x):
        x = torch.flatten(x, start_dim=1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return x

class decoder(nn.Module):

    def __init__(self, z_dim, h_dim1, h_dim2, img_height):

        super(decoder, self).__init__()

        self.img_height = img_height
        self.fc4 = nn.Linear(z_dim, h_dim2)
        self.fc5 = nn.Linear(h_dim2, h_dim1)
        self.fc6 = nn.Linear(h_dim1, img_height*img_height)

        
    def forward(self, x):
        x = F.relu(self.fc4(x))
        x = F.relu(self.fc5(x))
        x = F.sigmoid(self.fc6(x))
        x = x.view(-1, 1, self.img_height, self.img_height)
        return x



class VAE(pl.LightningModule):

    """
    Standard VAE with Gaussian Prior and approx posterior.
    Model is available pretrained on different datasets:
    """
    
    def __init__(
        self,
        img_height:int = 28,
        h_dim1:int = 512,
        h_dim2:int = 256,
        latent_dim:int = 2,
        kl_coeff: float = 0.1,
        lr: float = 1e-4,
        **kwargs
    ):
        """
        Args:
            input_height: height of the images
            enc_type: option between resnet18 or resnet50
            enc_out_dim: set according to the out_channel count of
                encoder used (512 for resnet18, 2048 for resnet50)
            kl_coeff: coefficient for kl term of the loss
            latent_dim: dim of latent space
            lr: learning rate for Adam
        """

        super(VAE, self).__init__()

        self.save_hyperparameters()

        self.lr = lr
        self.kl_coeff = kl_coeff
        self.latent_dim = latent_dim
        self.h_dim1 = h_dim1
        self.h_dim2 = h_dim2
        self.img_height = img_height

        self.fc_mu = nn.Linear(self.h_dim2, self.latent_dim)
        self.fc_var = nn.Linear(self.h_dim2, self.latent_dim)

        self.encoder = encoder(self.img_height, self.h_dim1, self.h_dim2)
        self.decoder = decoder(self.latent_dim, self.h_dim1, self.h_dim2, self.img_height)
</code></pre>
<p>The following training code is the same as the one in pl_bolts.models.autencoders.VAE. This gives me the following results (on MNIST):<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/2/5/252b0110c7889bc16c29cf1f4e15083b55cfdd46.png" alt="26_mnist" data-base62-sha1="5iNNKqeRyhje7oAhDcJQEW5j9Ma" width="302" height="392"><br>
(on my dataset):<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/a/cabaf6e32e4f78330193eedb21b422ee2564e223.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/cabaf6e32e4f78330193eedb21b422ee2564e223" title="49_chromo"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/cabaf6e32e4f78330193eedb21b422ee2564e223_2_384x500.jpeg" alt="49_chromo" data-base62-sha1="sVr18NJCtCMuHTkxQSGJxRQzXNh" width="384" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/cabaf6e32e4f78330193eedb21b422ee2564e223_2_384x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/cabaf6e32e4f78330193eedb21b422ee2564e223_2_576x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/c/a/cabaf6e32e4f78330193eedb21b422ee2564e223.jpeg 2x" data-dominant-color="0F0F0F"><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">49_chromo</span><span class="informations">662×860 50.4 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>I also tried another (smaller) convolutional architecture (with different configurations, like transposed convolutions, upsample then convolution, binary cross entropy or L1 loss) but I either get those results with no sample diversity, or the training collapses completely (completely white / black images). It really seems I’m missing something with the convolutional architectures.</p>
<p>So if anybody is aware of any gross mistake in this code, or a well known trick that needs to be followed to train convolutional VAEs, I’m all ears <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>