0
$\begingroup$

I am new to training VAEs and I am using it on some 16x16 images, that contains some images from a physics experiment with one or 2 events i.e. the images are mainly black, except for one or 2 regions where we have a Gaussian-like signal around several pixels. Below are the main components of my code:

class View(nn.Module):
    def __init__(self, size):
        super(View, self).__init__()
        self.size = size

    def forward(self, tensor):
        return tensor.view(self.size)

class Encoder(nn.Module):
    def __init__(self, lat_dim):
        super().__init__()

        N = 32
        
        self.encoder = nn.Sequential(
            nn.Conv2d(1, N, 4, 2, 1),          # B, N, 8, 8
            nn.ReLU(True),
            nn.Conv2d(N, N, 4, 2, 1),          # B,  N, 4, 4
            nn.ReLU(True),
            nn.Conv2d(N, 2*N, 4, 1),          # B,  N, 4, 4
            nn.ReLU(True),
            View((-1, 2*N*1*1)),                 # B, 2*N
        )    

        self.mu = nn.Linear(2*N, lat_dim)
        self.var = nn.Linear(2*N, lat_dim)

    def forward(self, x):
        
        x = self.encoder(x)
        z_mu = self.mu(x)
        z_var = self.var(x)
        
        return z_mu, z_var

class Decoder(nn.Module):
    def __init__(self, lat_dim):
        super().__init__()
        
        N = 32

        self.decoder = nn.Sequential(
            nn.Linear(lat_dim, 2*N),               # B, 256
            View((-1, 2*N, 1, 1)),               # B, 256,  1,  1
            nn.ReLU(True),
            nn.ConvTranspose2d(2*N, N, 4),      # B,  64,  4,  4
            nn.ReLU(True),
            nn.ConvTranspose2d(N, N, 4, 2, 1), # B,  64,  8,  8
            nn.ReLU(True),
            nn.ConvTranspose2d(N, 1, 4, 2, 1), # B,  32, 16, 16
        )
    
    def forward(self, x):
        x = self.decoder(x)

        return x

class VAE(nn.Module):
    def __init__(self, enc, dec):
        super().__init__()

        self.enc = enc
        self.dec = dec

    def forward(self, x):
        # encode
        z_mu, z_var = self.enc(x)

        # sample from the distribution having latent parameters z_mu, z_var
        # reparameterize
        std = torch.exp(z_var / 2)
        eps = torch.randn_like(std)
        x_sample = eps*std+z_mu 

        # decode
        predicted = self.dec(x_sample)
        return predicted, z_mu, z_var

lat_dim = 10 
model_E = Encoder(lat_dim) 
model_D = Decoder(lat_dim) 
model = VAE(model_E, model_D).cuda()

for epoch in range(801):
    model.train()
    
    if epoch%200==0 and epoch>1:
        lrs = lrs/10
    optimizer = optim.Adam(model.parameters(), lr = lrs, weight_decay=1e-3)
    
    for x, y in zip(train_data, train_data_labels):
        optimizer.zero_grad()

        x = x.float().cuda()
        y = y.float().cuda()
        
        x_sample, z_mu, z_var = model(x)        
        recon_loss = F.mse_loss(x_sample, x)
        kl_loss = torch.mean(-0.5 * torch.sum(1 + z_var - z_mu ** 2 - torch.exp(z_var), dim = 1), dim = 0)
        
        beta = 0.1
        loss = recon_loss + beta*kl_loss
        
        loss.backward()
        optimizer.step()

The data is normalized to mean 0 and std of 1. When I train it however, the KL divergence loss goes almost instantly (I guess it starts there) to ~10^-5 and keeps going down to around 10^-10, while the reconstruction loss doesn't change from the starting values (around 1). I tried different values for the beta parameter (the weight of the kl loss) and i have to set that to 10^-4 or lower for the reconstruction loss to actually go down. But if I do that the KL loss doesn't change much (which makes sense I guess). But in most papers I read about VAE the beta parameter is between 0 and 1, so I guess I am doing something wrong. Also my data is a lot simpler in both structure and size compared to MNIST (or other complex data sets), so I guess that a VAE code should work. What am I doing wrong? I am using a batch size of 128 and 6000 images for training. The loss mentioned in the post is the training loss. Thank you!

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.