解耦VAE

此仓库包含用于研究VAE中解耦现象的代码（训练、指标计算、绘图），并使用单一架构比较5种不同的损失函数（差异总结）：

• 来自《变分贝叶斯自动编码》的标准VAE损失
• 来自《β-VAE：利用约束变分框架学习基本视觉概念》的β-VAE_H
• 来自《理解β-VAE中的解耦》的β-VAE_B
• 来自《通过因子分解实现解耦》的FactorVAE
• 来自《在变分自编码器中分离解耦源》的β-TCVAE

注意事项：

• 已测试兼容Python 3.6及以上版本
• 已测试支持CPU和GPU

目录：

1. 安装
2. 运行
3. 绘图
4. 数据
5. 我们的贡献
6. 损失函数解释
7. 引用

安装

# 克隆仓库
pip install -r requirements.txt

运行

使用python main.py <模型名称> <参数>来训练和/或评估模型。例如：

python main.py btcvae_celeba_mini -d celeba -l btcvae --lr 0.001 -b 256 -e 5

您也可以使用-x <实验>来运行预定义的实验及超参数，这些超参数可在hyperparam.ini中找到。每个实验的预训练模型可从results/<实验>目录中获取（通过./bin/train_all.sh脚本生成）。

输出

这将创建一个名为results/<保存名称>/的目录，其中包含：

• model.pt：训练结束时的模型。
• model-i.pt：每经过i次迭代后保存的模型检查点，默认每10次保存一次。
• specs.json：运行程序时使用的参数（默认参数及通过命令行修改后的参数）。
• training.gif：训练过程中每一epoch的潜在空间维度Z的遍历动画GIF。
• train_losses.log：训练期间计算的所有（子）损失。
• test_losses.log：训练结束后以评估模式运行时计算的所有（子）损失（不进行采样）。
• metrics.log：互信息差距指标和轴对齐指标。仅当指定--is-metric时才会计算（较慢）。

帮助

用法：main.py ...

PyTorch实现与评估解耦变分自编码器及其指标。
可选参数：
  -h, --help            显示此帮助信息并退出

通用选项：
  name                  模型名称，用于存储或加载。
  -L, --log-level {CRITICAL,ERROR,WARNING,INFO,DEBUG,NOTSET}
                        日志级别。（默认：info）
  --no-progress-bar     禁用进度条。（默认：False）
  --no-cuda             即使有CUDA设备也禁用CUDA训练。（默认：False）
  -s, --seed SEED       随机种子。可以设置为`None`以获得随机行为。（默认：1234）

训练特定选项：
  --checkpoint-every CHECKPOINT_EVERY
                        每隔n个epoch保存一次训练好的模型检查点。（默认：30）
  -d, --dataset {mnist,fashion,dsprites,celeba,chairs}
                        训练数据路径。（默认：mnist）
  -x, --experiment {custom,debug,best_celeba,VAE_mnist,VAE_fashion,VAE_dsprites,VAE_celeba,VAE_chairs,betaH_mnist,betaH_fashion,betaH_dsprites,betaH_celeba,betaH_chairs,betaB_mnist,betaB_fashion,betaB_dsprites,betaB_celeba,betaB_chairs,factor_mnist,factor_fashion,factor_dsprites,factor_celeba,factor_chairs,btcvae_mnist,btcvae_fashion,btcvae_dsprites,btcvae_celeba,btcvae_chairs}
                        预定义的实验任务。若非`custom`，则会覆盖其他部分参数。（默认：custom）
  -e, --epochs EPOCHS   最大训练轮数。（默认：100）
  -b, --batch-size BATCH_SIZE
                        训练批次大小。（默认：64）
  --lr LR               学习率。（默认：0.0005）

模型特定选项：
  -m, --model-type {Burgess}
                        使用的编码器和解码器类型。（默认：Burgess）
  -z, --latent-dim LATENT_DIM
                        潜在变量的维度。（默认：10）
  -l, --loss {VAE,betaH,betaB,factor,btcvae}
                        使用的VAE损失函数类型。（默认：betaB）
  -r, --rec-dist {bernoulli,laplace,gaussian}
                        每个像素所使用的似然分布形式。（默认：伯努利）
  -a, --reg-anneal REG_ANNEAL
                        渐进式添加正则化项的退火步数。具体退火内容因损失而异。（默认：0）

BetaH特有参数：
  --betaH-B BETAH_B     KL散度项的权重（论文中的beta）。（默认：4）

BetaB特有参数：
  --betaB-initC BETAB_INITC
                        初始退火容量。（默认：0）
  --betaB-finC BETAB_FINC
                        最终退火容量。（默认：25）
  --betaB-G BETAB_G     KL散度项的权重（论文中的gamma）。（默认：1000）

factor VAE特有参数：
  --factor-G FACTOR_G   TC项的权重（论文中的gamma）。（默认：6）
  --lr-disc LR_DISC     判别器的学习率。（默认：5e-05）

beta-tcvae特有参数：
  --btcvae-A BTCVAE_A   MI项的权重（论文中的alpha）。（默认：1）
  --btcvae-G BTCVAE_G   维度级KL项的权重（论文中的gamma）。（默认：1）
  --btcvae-B BTCVAE_B   TC项的权重（论文中的beta）。（默认：6）

评估特定选项：
  --is-eval-only        是否仅使用预先计算好的模型`name`进行评估。（默认：False）
  --is-metrics          是否计算解耦指标。目前仅适用于`dsprites`数据集，因为它是唯一已知真实变化因素的数据集。（默认：False）
  --no-test             是否不计算测试损失。（默认：False）
  --eval-batchsize EVAL_BATCHSIZE
                        评估时的批次大小。（默认：1000）

绘图

使用 python main_viz.py <model-name> <plot_types> <param> 可以利用预训练模型进行绘图。例如：

python main_viz.py btcvae_celeba_mini gif-traversals reconstruct-
                        traverse -c 7 -r 6 -t 2 --is-posterior

这会将图表保存在模型目录 results/<model-name>/ 中。所有实验生成的图表都位于各自的目录中（通过 ./bin/plot_all.sh 脚本创建）。

帮助信息

用法: main_viz.py ...

用于使用 `disvae` 的预训练模型进行绘图的命令行界面

位置参数:
  name                  模型名称，用于存储和加载。
  {generate-samples,data-samples,reconstruct,traversals,reconstruct-traverse,gif-traversals,all}
                        需要生成的所有图表类型。`generate-samples`: 随机解码样本。`data-samples`: 数据集中的样本。`reconstruct`: 前半部分为原始样本，后半部分为对应的重建样本。`traversals`: 沿着最重要的几个潜在维度进行遍历，每列展示来自先验或后验分布的不同样本。`reconstruct-traverse`: 第一行是原始样本，第二行是重建样本，其余行是遍历结果。`gif-traversals`: 由多张动图组成的网格，行代表潜在维度，列代表示例，每张动图展示后验分布的遍历过程。`all`: 运行所有类型的图表。

可选参数:
  -h, --help            显示帮助信息并退出
  -s, --seed SEED       随机种子。可以设置为 `None` 以获得随机行为。（默认值：None）
  -r, --n-rows N_ROWS   可视化时使用的行数（如果适用）。（默认值：6）
  -c, --n-cols N_COLS   可视化时使用的列数（如果适用）。（默认值：7）
  -t, --max-traversal MAX_TRAVERSAL
                        潜在变量遍历时的最大位移量。假设遍历是对称的。若 `m>=0.5`，则采用绝对值遍历；若 `m<0.5`，则基于分布的百分比（分位数）进行遍历。例如，在先验分布中，分布为标准正态分布，因此 `m=0.45` 对应于绝对值 `1.645`，因为 `2m=90%` 的标准正态分布在 `-1.645` 到 `1.645` 之间。需要注意的是，在后验分布中，分布不再是标准正态分布。（默认值：2）
  -i, --idcs IDCS [IDCS ...]
                        放置在样本开头的图像索引列表。（默认值：空列表）
  -u, --upsample-factor UPSAMPLE_FACTOR
                        图像上采样的缩放因子（如果适用）。（默认值：1）
  --is-show-loss        在图表中显示损失值（如果适用）。（默认值：False）
  --is-posterior        使用后验分布而非先验分布进行遍历。（默认值：False）

示例

以下是一些你可以生成的图表示例：

• python main_viz.py <model> reconstruct-traverse --is-show-loss --is-posterior 第一行是原始样本，第二行是重建样本，其余行是潜在变量遍历结果。以 btcvae_dsprites 为例：

  

• python main_viz.py <model> gif-traversals 由多张动图组成的网格，行代表潜在维度，列代表示例，每张动图展示后验分布的遍历过程。以 btcvae_celeba 为例：

  

• 使用 bin/plot_all.sh 脚本生成的动图网格。网格的列对应不同的数据集（除 FashionMNIST 外），行对应不同的模型（顺序为：标准 VAE、β-VAE_H、β-VAE_B、FactorVAE、β-TCVAE）：

  

更多示例，请参阅预先定义的实验所生成的所有图表，它们都存放在各自的目录中（通过 ./bin/plot_all.sh 脚本创建）。

数据

当前可用的数据集包括：

• MNIST
• FashionMNIST
• 3D Chairs
• Celeba
• 2D Shapes / Dsprites

首次运行时，数据集将会被下载并存储在 data 目录下，供后续使用。下载可能需要一些时间，且如果下载链接发生变化，可能会导致无法继续下载。在这种情况下，您可以：

1. 提交一个问题报告。
2. 在 utils/datasets.py 文件中修改您所需数据集的 URL（请记得提交一个 Pull Request :) ）。
3. 手动下载数据并以相同文件名保存（不推荐）。

我们的贡献

除了复现上述论文之外，我们还提出并研究了以下内容：

轴对齐度量指标

由于定性评估具有主观性和耗时性，因此不适合用来可靠地比较不同模型。近期的一些论文采用了基于真实变化因素 v 和潜在维度 z 的定量解耦度量方法。其中，互信息差距 (MIG) 是一种颇具吸引力的信息论度量指标，因为它无需使用任何分类器。在 dSprites 数据集中，我们有 10 个潜在维度和 5 个生成因素，若要使 MIG 达到 1，就需要让其中 5 个潜在维度精确地编码这些真实的变化因素，而其余 5 个则与前 5 个无关。

尽管长期来看我们希望使用类似 MIG 的指标，但目前大多数模型的得分并不理想，且难以明确其改进方向。因此，我们提出了轴对齐度量指标 AAM，该指标并不关注 z 编码了多少 v 的信息，而是关注每个 v_k 是否仅由单个 z_j 编码。例如，在 dSprites 数据集中，即使 z 只编码了形状 x 位置方差的 90%，只要这 90% 完全由单一潜在维度 z_j 编码，就可以达到 AAM = 1。这一指标有助于更好地理解各个模型的优势和不足。形式化定义如下：

其中下标 (d) 表示第 d 个次序统计量，而 I_x 则通过经验分布和分层抽样来估计（类似于 MIG）：

单一模型对比

该模型与所有损失函数解耦，因此在不修改损失函数的情况下，很容易对编码器/解码器进行修改。我们仅使用单一模型，以便更客观地比较不同的损失函数。所使用的模型来自论文《理解 β-VAE 中的解耦》，其架构如下所示：

损失函数说明

所有先前的损失函数都是以下损失函数的特例：

1. 索引-代码互信息：潜在变量 z 与数据变量 x 之间的互信息。关于如何正确处理该项，学术界存在争议。从信息瓶颈的角度来看，这一项应当受到惩罚。而 InfoGAN 则通过增加互信息（负 α）取得了良好效果。最后，Wasserstein 自编码器 直接舍弃了该项。

2. 总相关性（TC）：潜在变量的联合分布与其边缘分布乘积之间的 KL 散度。即潜在空间各维度之间依赖性的度量。增大 β 值会迫使模型在数据分布中寻找统计独立的变化因子。

3. 逐维 KL 散度：后验分布的每个维度与先验分布之间的 KL 散度。该项确保学习到一个接近先验分布的紧凑空间，从而能够采样生成新的样本。

不同损失函数在估算上述各项以及所使用的超参数方面存在差异：

• 标准 VAE 损失：α=β=ɣ=1。各项均通过闭式解精确计算（先验与后验之间的 KL 散度）。为最紧的下界。
• β-VAE_H：α=β=ɣ>1。各项同样通过闭式解精确计算。只是在 KL 散度前增加了一个超参数（文中为 β）。
• β-VAE_B：α=β=ɣ>1。与 β-VAE_H 类似，但仅当三项偏离随训练过程逐渐增大的容量 C 时才予以惩罚。
• FactorVAE：α=ɣ=1，β>1。各项均通过闭式解精确计算。同样在 KL 散度前增加了一个超参数（文中为 β），并在标准 VAE 损失基础上添加了一个加权的总相关性项。总相关性采用分类器和密度比技巧进行估计。需要注意的是，他们论文中的 ɣ 在我们的框架中对应于 β+1。
• β-TCVAE：α=ɣ=1（尽管可以调整），β>1。概念上等同于 FactorVAE，但各项分别使用小批量分层采样进行估计。

引用

在学术研究中使用本仓库实现的任一模型时，请引用相应的论文（链接位于 README 的顶部）。若需引用此特定实现，可使用以下格式：

@misc{dubois2019dvae,
  title        = {Disentangling VAE},
  author       = {Dubois, Yann and Kastanos, Alexandros and Lines, Dave and Melman, Bart},
  month        = {march},
  year         = {2019},
  howpublished = {\url{http://github.com/YannDubs/disentangling-vae/}}
}

disentangling-vae 快速上手指南

disentangling-vae 是一个基于 PyTorch 的开源项目，用于研究和比较变分自编码器（VAE）中的解耦（Disentanglement）能力。它实现了标准 VAE、$\beta$-VAE、FactorVAE 和 $\beta$-TCVAE 等五种主流损失函数，并提供了完整的训练、评估及可视化工具。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows
• Python 版本：>= 3.6 (推荐 3.8+)
• 硬件支持：支持 CPU 和 GPU (CUDA) 训练
• 核心依赖：PyTorch, NumPy, Matplotlib 等（将通过 requirements.txt 自动安装）

国内加速建议：

1. 建议使用国内镜像源安装 Python 依赖，例如清华源或阿里源，以加快下载速度。
2. 项目首次运行时会自动下载数据集（如 MNIST, CelebA 等）。若官方源下载失败，可手动下载数据至 data 目录，或修改 utils/datasets.py 中的下载链接。

安装步骤

1. 克隆仓库

   git clone https://github.com/YannDubs/disentangling-vae.git
   cd disentangling-vae
   

2. 安装依赖 推荐使用国内镜像源安装依赖包：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

1. 训练模型

使用 main.py 脚本进行模型训练或评估。您可以指定模型类型、数据集、损失函数及超参数。

示例命令： 以下命令使用 $\beta$-TCVAE 损失函数在缩小的 CelebA 数据集上训练 5 个 epoch：

python main.py btcvae_celeba_mini -d celeba -l btcvae --lr 0.001 -b 256 -e 5

常用参数说明：

• -d: 数据集选择 (mnist, fashion, dsprites, celeba, chairs)
• -l: 损失函数类型 (VAE, betaH, betaB, factor, btcvae)
• -e: 训练轮数 (epochs)
• -b: 批次大小 (batch size)
• --lr: 学习率
• -x: 使用预定义实验配置（覆盖其他参数），例如 -x best_celeba

训练完成后，结果将保存在 results/<model-name>/ 目录下，包含模型权重 (model.pt)、训练日志及潜变量遍历的 GIF 动图。

2. 可视化结果

使用 main_viz.py 脚本对预训练模型进行可视化分析，包括重构图像、潜变量遍历等。

示例命令： 生成潜变量遍历的 GIF 动图及重构对比图：

python main_viz.py btcvae_celeba_mini gif-traversals reconstruct-traverse -c 7 -r 6 -t 2 --is-posterior

常用参数说明：

• plot_types: 绘图类型，可选 generate-samples, reconstruct, traversals, gif-traversals, all 等。
• -r / -c: 设置输出图像的行数和列数。
• --is-posterior: 基于后验分布进行遍历（默认为先验分布）。
• --is-show-loss: 在图中显示损失值。

生成的图片将保存至对应的 results/<model-name>/ 目录中。