I am trying to train a VAE on shapenet but I can't seem to make it work. Any help or ideas would be highly appreciated. Now the problem is whenever I apply the KL divergence loss the network seems to get stuck with having a high KL divergence and high reconstruction loss. Even though without the KL divergence (plain auto encoder) the reconstructions are almost perfect.
Problem Setup:
Input: 64^3 TSDFs of shapenet
For now, I am just trying with the chair category of around (6.5k chairs)
Ouput: 64^3 TSDF- reconstruction
Network Architecture: I am mostly using the same architecture as in AutoSDF, which is a fully 3D convolutional encoder and decoder to output the final TSDF matrix. I tried with various latent space sizes but none seems to work.
KL divergence loss function:
loss = -0.5 * torch.sum(1 + logvar.flatten(1) - mu.flatten(1).pow(2) -
logvar.flatten(1).exp(), dim=1)
loss = torch.mean(loss, dim=0)
where mu and logvar are the 3D outputs of the encoder (tried with different latent space sizes as explained below).
List of experiment settings
- Latent spaces of
(2*8*8*8),(4*8*8*8*8),(8*8*8*8*8)(16*8*8*8*8),(32*8*8*8)and(64*8*8*8): None of them seemed to work in the VAE setting, but(16*8*8*8*8),(32*8*8*8)and(64*8*8*8)worked well in the regular AE setting. - Added linear layers to the final output to obtain
muandlogvarwith 512 dimensionality in one experiment and 1024 in another -> made things worse - Followed all suggestions in this thread. Including "KL cost annealing" and "cyclic annealing".
- Tried to normalize the KL loss weighting by the dimensionality of the latent space
- Several combinations of reconstruction weighting and KL divergence weighting
- Used the sum reduction instead of mean to get a higher priority for the reconstruction loss
None of the above seems to work an I have done a fair bit of research now, but I have no idea what to do next so I would appreciate any help or discussion.
