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!