murphyk
commented
2 years ago It would also be helpful to make a separate notebook that compares AE and VAE on MNIST so we can reproduce figs 20.24 and 20.25 from book1 (see below). You can factor out common code and put it in https://github.com/probml/probml-utils/.
Currently the reconstruction accuracy of the AE looks bad, which makes me suspect there is something wrong. Also the current code relies on pytorch code in https://github.com/probml/probml-utils/blob/main/probml_utils/ae_mnist_conv.py and https://github.com/probml/probml-utils/blob/main/probml_utils/vae_conv_mnist.py, which I would like to deprecate.
nalzok
figs 21.3 and 21.4 in book 2 (shown below) compare the reconstructon and unconditional sampling abilities of a convolutional AE and VAE. However these figures were generated by https://github.com/probml/pyprobml/blob/master/notebooks/book2/20/vae_compare_results.ipynb which relies on this somewhat complex pytorch lightning library: https://github.com/probml/pyprobml/tree/master/vae.
Please reproduce this basic demo using code that is similar to https://github.com/probml/pyprobml/blob/master/notebooks/book1/20/ae_mnist_conv_jax.ipynb. To get the celeba data you will have to follow the instructions at https://github.com/probml/pyprobml/tree/master/vae#download-celeba-data. For simplicity it is probably best to make a single colab that does everything, assuning that training can be done in under 5 minutes (if not, we may need to pretrain the models and store the checkpoints).