Open mmubeen-6 opened 3 years ago
Hi @mmubeen-6, any luck on reproducing this work? I re-implemented the model from the ground up and have yet to generate comparable results relative to the original paper. I'll try a few more hyperparameter settings and train longer than the original paper (which may help?), but I'd be interested to hear if you've solved any of your pre-existing issues.
It's been a long time since I ran those experiments, but a few things that will improve the generated images are:
It's also worth noting that the encoder in spatial-VAE can sometimes get stuck in bad local optima (especially regarding rotation inference) which then leads to bad generator performance. There are a few tricks implemented here to try to avoid those, but they only work so-so (e.g., including rotated images as input, but decoding the unrotated image by shifted the predicted rotation by the known augmentation rotation). You might want to take a look at some newer work (https://arxiv.org/abs/2210.12918, https://github.com/SMLC-NYSBC/TARGET-VAE) where we improved the encoder to address some of these issues. We lightly tested that on galaxy zoo but didn't push it as far as it should be able to go with a larger spatial generator and/or better initial featurization of the coordinates
Hi @tbepler, thanks for your comments! I spent the evening tinkering with a bit of my code and I did notice a bit of the local optima issues on a few runs, so I will play with it a bit more taking your suggestions into account. I also did see your TARGET-VAE paper pop up as well too, congrats on that! I will try playing around with the group convolutions if I get a chance.
If you don't mind, I might create another github repository with a bit of a refactored/re-engineered version of the spatial-VAE (with proper attributions and references of course). I think this could be a nice architecture for some of the applied bio stuff I'm working on - so I will probably run some additional experiments with more expressive encoders, different ways to condition the decoding with the latent variables, cyclical annealing of the kld, and maybe swapping out linear layers for 1x1 convolutions, etc. etc. I may potentially add in reflection if I have time.
I also wonder if semi-supervised learning can help with convergence; since it's seems reasonably straight forward to aggregate some ground truth rotations/translations either through augmentation or extracting them via fits (e.g. major axis or something).
@tomouellette you're welcome to fork the code and use it however you like. I agree semi-supervised learning could help with convergence if you have labeled angles for a subset of images. It should be pretty straightforward to include in the objective.
HI @tbepler, I am trying to reproduce the results of your paper on the galaxy dataset but unable to exactly achieve those. Could you please share the exact training parameters. I am currently using the following command to train it.
python3 train_galaxy.py galaxy_zoo/galaxy_zoo_train.npy galaxy_zoo/galaxy_zoo_test.npy -d 0 --num-epochs 300 --save-prefix galaxy_zoo_models/testing -z 100 --minibatch-size 100 --dx-scale 0.125
.Moreover, in order visualize the reconstructed images, I am using the following code snippet. Please have a look at it.