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_