Hello,
I’m training a VAE and I’m a bit confused as to what I need to do during testing.
Based on what I’ve read, it seems that during training we use the reparam trick (as per the Kingma paper 2) and then during testing we simply sample from epsilon, ie N(0,1). This is how I’m trying to implement this:
    def train_reparam(self, h):
        mean, logvar = self.fc1a(h), self.fc1b(h)  
        std_dev = torch.exp(0.5*logvar)
        eps = torch.randn_like(std_dev)        
        z = mean + std_dev * eps
        return z, mean, logvar
    
    def test_reparam(self, h):
        mean, logvar = self.fc1a(h), self.fc1b(h)  
        std_dev = torch.exp(0.5*logvar)
        eps = torch.randn_like(std_dev)        
        z = eps # z is now eps (norm dist)
        return z, mean, logvar

    def forward(self, x):
        h = self.encoder(x)
        if self.training:
            z, mean, logvar = self.train_reparam(h)
        else:
            z, mean, logvar = self.test_reparam(h)
        z = self.fc2(z)
        return self.decoder(z), mean, logvar

The results obtained with the above, however, as confusing. The test loss seems to be lower than during training and the quality of images generated during testing is poor.
What also confuses me is that when I look for examples of VAE code, many of them don’t seem to operate differently during testing (ie they reparam at train and test) - even the PyTorch VAE example 4 !
So my questions are:

Is what I’m doing correct? Do I need this train/test distinction?
If so, should I return z, mean & logvar during test or should that be replaced with z, 0, 1?

Any help would be greatly appreciated!
Thanks,
NK