CVAE与VQ-VAE实现

这是一个实现向量量化变分自编码器（VQ-VAE）和卷积变分自编码器（CVAE）的示例，用于压缩MNIST和CIFAR-10数据集。 基于论文《神经离散表示学习》(Neural Discrete Representation Learning)[1]。 代码构建于PyTorch的VAE示例之上。

pip install -r requirements.txt
python main.py

系统要求

• Python 3.6（可能Python 3.5也适用）
• PyTorch 0.4
• 其他具体需求请参考requirements.txt文件

使用方法

# 示例用法
python3 main.py --dataset=cifar10 --model=vqvae --data-dir=~/.datasets --epochs=3

结果展示

所有图像均来源于测试集。上排为原始图像，下排为重构后的图像。

其中，k代表词典中的元素数量，d代表词典中元素的维度（瓶颈层的通道数）。

• MNIST (k=10, d=64)

• CIFAR-10 (k=128, d=256)

• ImageNet (k=512, d=128)

待办事项:

• 实现《离散潜变量图像模型的连续放松训练》[2]
• 使用PixelCNN先验进行采样
• 提高在CIFAR-10上的表现——字典查找应从1扩展到10个近邻
• 提高在CIFAR-10上的表现——将MSE损失替换为NLL损失
• 提高在CIFAR-10上的表现——计算每维比特数
• 将模型架构与官方实现对比
• 将MNIST和Cifar的VAE与VQ-VAE合并至单一脚本

致谢

感谢tf-vqvae项目作为一个良好的参考资源。

[1]：van den Oord, A., et al. "Neural discrete representation learning." arXiv preprint arXiv:1711.00937 (2017). [2]: Exploring Continuous Relaxations of Discrete Variables for Latent Variable Models, Bayesian Deep Learning Workshop, NIPS 2017.