vdvae-jax:JAXFlax中非常深的VAE

JAX /亚麻中非常深的VAE 从移植而来的使用和。 我试图使该实现尽可能接近原始实现。 我能够重用大部分代码，包括仍然使用PyTorch的数据输入管道。 我建议为此安装PyTorch的仅CPU版本。 使用JAX 0.2.10，亚麻0.3.0，PyTorch 1.7.1，NumPy 1.19.2进行了测试。 我还对cifar10进行了收敛训练，并使用--conv_precision=highest从本文中复制了测试ELBO值2.87，请参见下文。 如果有人要求为cifar训练有素的检查站，我将很乐意上传它们。 从本文中可以得出一些模型样本及其可视化方式： 设置 除了JAX，Flax，NumPy和PyTorch之外，此实现还取决于和 ： pip install pillow pip install sklearn 另外，您还必须下载数据，具体取决于要运行的数据： ./setup_