disentanglement_lib

disentanglement_lib 是一个用于研究学习解纠缠表征的开源库。 它支持多种不同的模型、指标和数据集：

• 模型: BetaVAE、FactorVAE、BetaTCVAE、DIP-VAE
• 指标: BetaVAE 分数、FactorVAE 分数、互信息差距、SAP 分数、DCI、MCE、IRS、UDR
• 数据集: dSprites、Color/Noisy/Scream-dSprites、SmallNORB、Cars3D 和 Shapes3D
• 它还包含 10,800 个预训练的解纠缠模型（详情见下文）。

disentanglement_lib 由谷歌大脑苏黎世的 Olivier Bachem 和 Francesco Locatello 创建，用于大规模实证研究

挑战解纠缠表征无监督学习中的常见假设。 Francesco Locatello、Stefan Bauer、Mario Lucic、Gunnar Rätsch、Sylvain Gelly、Bernhard Schölkopf、Olivier Bachem。ICML（最佳论文奖），2019 年。

该代码使用 Python 3 测试，旨在运行在 Linux 系统上（例如 Google Cloud 深度学习虚拟机)。 它使用 TensorFlow、Scipy、Numpy、Scikit-Learn、TFHub 和 Gin。

它是如何工作的？

disentanglement_lib 包含多个不同的步骤：

• 模型训练: 训练一个 TensorFlow 模型，并将训练好的模型保存为 TFHub 模块。
• 后处理: 获取一个训练好的模型，提取表征（例如通过使用高斯编码器的均值），并将表征函数保存为 TFHub 模块。
• 评估: 获取一个表征函数并计算解纠缠指标。
• 可视化: 获取一个训练好的模型并进行可视化。

所有配置细节和不同步骤的实验结果都会被保存并沿步骤传播（详见下文）。最后，它们可以汇总到一个 JSON 文件中，并用 Pandas 进行分析。

使用方法

安装 disentanglement_lib

首先，克隆此仓库：

git clone https://github.com/google-research/disentanglement_lib.git

然后，进入仓库目录（使用 cd disentanglement_lib），运行：

pip install .[tf_gpu]

（或 pip install .[tf] 以安装不带 GPU 支持的 TensorFlow）。 这将安装该包及所有必需的依赖项。 要验证一切是否正常工作，只需运行测试套件：

dlib_tests

下载数据集

要下载训练模型所需的数据，进入任意文件夹并运行：

dlib_download_data

这将在当前工作目录中安装所有必需的数据文件（Shapes3D 数据集除外，因为它并未公开发布）。 为方便起见，我们建议将环境变量 DISENTANGLEMENT_LIB_DATA 设置为此路径，例如在你的 .bashrc 文件中添加：

export DISENTANGLEMENT_LIB_DATA=<数据目录路径>

如果你选择不设置环境变量 DISENTANGLEMENT_LIB_DATA，disentanglement_lib 将始终在当前文件夹中查找数据。

复现先前实验

要完整训练并评估论文《挑战解纠缠表征无监督学习中的常见假设》中 12,600 个模型之一，只需运行：

dlib_reproduce --model_num=<?>

其中 <?> 应替换为 0 到 12,599 之间的模型索引，对应于要训练的模型 ID。 这将花费几个小时，并在 output/<?> 文件夹中生成训练好的模型（包括检查点和 TFHub 模块）、实验结果（JSON 格式）以及可视化内容（包括 GIF 动画）。 如果只想打印该模型的配置而不进行训练，可添加标志 --only_print。

在训练完若干模型后，你可以通过运行以下命令汇总结果（在同一文件夹中）：

dlib_aggregate_results

这将创建一个 results.json 文件，包含所有汇总的结果。

运行不同配置

disentanglement_lib 内部使用 gin 来配置超参数和其他设置。 要训练提供的模型之一但使用不同的超参数，你需要编写一个 gin 配置文件，例如 examples/model.gin。 然后，你可以使用以下命令：

dlib_train --gin_config=examples/model.gin --model_dir=<模型输出目录>

来训练模型，其中 --model_dir 指定保存结果的路径。

要按照论文《挑战解纠缠表征无监督学习中的常见假设》中的评估协议评估新训练的模型，只需运行：

dlib_reproduce --model_dir=<模型输出目录> --output_directory=<输出目录>

同样，如果你想自定义表征的提取和评估方式，也可以查看 dlib_postprocess 和 dlib_evaluate。

开始你自己的研究

disentanglement_lib 易于扩展，可用于实现与解纠缠表征相关的新型模型和指标。 要开始，只需浏览 examples/example.py，它向你展示了如何创建自己的解纠缠模型和指标，以及如何将其与现有模型和指标进行基准测试。

预训练的disentanglement_lib模块

重现研究挑战无监督学习中解纠缠表征的常见假设中的全部12,600个模型需要大量的计算资源。 为了促进进一步的研究，disentanglement_lib包含10,800个预训练的disentanglement_lib模块，这些模块对应于使用dlib_reproduce并指定--model_num=<?>在0到10,799之间运行的结果（其他模型对应于未公开发布的Shapes3D）。 每个disentanglement_lib模块包含训练好的模型（以TFHub模块的形式）、提取的表征（同样以TFHub模块的形式），以及记录的实验结果，例如不同的解纠缠分数（以JSON格式存储）。 这使得新模型与预训练模型的比较变得简单，并且能够在预训练模型集合上计算新的解纠缠指标。

要访问这10,800个预训练的disentanglement_lib模块，您可以使用以下链接下载单个模块：

https://storage.googleapis.com/disentanglement_lib/unsupervised_study_v1/<?>.zip

其中<?>对应于0到10,799之间的模型索引（示例）。

存储桶中的每个ZIP文件对应于使用该模型编号运行一次dlib_reproduce的结果。 要了解更多关于所用配置设置的信息，请查看disentanglement_lib/config/unsupervised_study_v1/sweep.py中的代码，或者运行：

dlib_reproduce --model_num=<?> --only_print

常见问题

如何为我的模型制作精美的GIF动画？

如果您运行dlib_reproduce，它们会自动保存到输出目录的visualizations子文件夹中。否则，您可以使用脚本dlib_visualize_dataset生成这些动画，或调用disentanglement_lib/visualize/visualize_model.py中的函数visualize(...)。

结果和模型是如何保存的？

在每个主要步骤（训练/后处理/评估）完成后，都会创建一个输出目录。 对于所有步骤，都有一个results文件夹，其中包含该步骤之前的所有配置设置和实验结果。 gin子文件夹包含以gin格式存储的每一步的运行配置。 json子文件夹包含以JSON格式存储的运行配置和该步骤的实验结果。 最后，aggregate子文件夹包含聚合的JSON文件，每个文件同时包含所有先前步骤的配置和结果。

训练步骤还会额外保存TensorFlow检查点（存放在tf_checkpoint子文件夹中）和训练好的模型作为TFHub模块（存放在tfhub子文件夹中）。类似地，后处理步骤会将表征函数保存为TFHub模块（存放在tfhub子文件夹中）。 如果您运行dlib_reproduce，它会为所有您运行过的不同子步骤创建子文件夹。特别是，它会为每个您计算的指标创建一个输出目录。

如何访问结果？

要访问结果，首先使用dlib_aggregate_results聚合所有结果，指定一个能匹配所有结果文件的glob模式。 例如，在使用dlib_reproduce训练几个不同模型后，您可以指定：

dlib_aggregate --output_path=<...>.json \
  --result_file_pattern=<...>/*/metrics/*/*/results/aggregate/evaluation.json

glob模式中的第一个会匹配不同的模型，第二个匹配不同的表征，最后一个*匹配不同的指标。 最后，您可以这样访问聚合后的结果：

from disentanglement_lib.utils import aggregate_results
df = aggregate_results.load_aggregated_json_results(output_path)

在代码中应该查找什么？

以下是对整体代码结构的指南：

(1) 训练步骤：

• disentanglement_lib/methods/unsupervised： 包含训练协议（train.py）以及所有用于训练方法的模型函数（vae.py）。这些方法都继承自GaussianEncoderModel类。
• disentanglement_lib/methods/shared： 包含不同模型中共享的架构、损失和优化器。

(2) 后处理步骤：

• disentanglement_lib/postprocess： 包含后处理流程（postprocess.py）以及两种提取方法（methods.py）。

(3) 评估步骤：

• disentanglement_lib/evaluation：包含评估协议（evaluate.py）。

• disentanglement_lib/evaluation/metrics：包含不同解纠缠指标的实现。

超参数和配置文件：

• disentanglement_lib/config/unsupervised_study_v1： 包含不同步骤的gin配置文件（*.gin），以及论文[挑战无监督学习中解纠缠表征的常见假设]中实验的超参数扫描（sweep.py）。

共享功能：

• bin：用于运行不同流程、可视化数据集和模型，以及聚合结果的脚本。

• disentanglement_lib/data/ground_truth： 包含生成数据的所有脚本。所有数据集（在named_data.py中）都是GroundTruthData类的实例。

• disentanglement_lib/utils： 包含帮助函数，用于聚合和保存流程结果以及训练好的模型。

• disentanglement_lib/visualize： 包含用于数据集和训练模型的可视化函数。

NeurIPS 2019解纠缠挑战

该库也用于NeurIPS 2019解纠缠挑战。挑战包含三个不同的数据集。

1. 简单的渲染图像（mpi3d_toy）
2. 真实的渲染图像（mpi3d_realistic）：尚未发布
3. 真实世界图像（mpi3d_real）：尚未发布

目前，只有简单的渲染数据集是公开可用的，运行以下命令即可自动下载：

dlib_download_data

其他数据集将在比赛的后期阶段提供。更多关于比赛的信息，请访问比赛网站。有关数据集的更多信息，请参阅这里或arXiv预印本关于从仿真到真实世界的归纳偏置迁移：一个新的解纠缠数据集。

抽象推理实验

该库还包含用于以下论文实验的代码，位于disentanglement_lib/evaluation/abstract_reasoning子目录中：

解纠缠表征对抽象视觉推理有帮助吗？ Sjoerd van Steenkiste、Francesco Locatello、Jürgen Schmidhuber、Olivier Bachem。NeurIPS，2019年。

实验方案分为两部分： 首先，要训练解纠缠模型，可以使用标准的复现流程（dlib_reproduce），例如通过以下命令：

dlib_reproduce --model_num=<?> --study=abstract_reasoning_study_v1

其中<?>应替换为0到359之间的模型索引， 对应于要训练的模型ID。

其次，要训练抽象推理模型，可以使用自动安装的流程dlib_reason。 要配置模型，请根据需要复制并修改disentanglement_lib/config/abstract_reasoning_study_v1/stage2/example.gin。 然后，使用以下命令来训练和评估一个抽象推理模型：

dlib_reason --gin_config=<?> --input_dir=<?> --output_dir=<?>

结果将保存在输出目录的results子目录中。

公平性实验

该库还包含用于以下论文实验的代码，位于disentanglement_lib/evaluation/metrics/fairness.py：

解纠缠表征的公平性研究 Francesco Locatello、Gabriele Abbati、Tom Rainforth、Stefan Bauer、Bernhard Schoelkopf、Olivier Bachem。NeurIPS，2019年。

要训练和评估所有模型，只需使用以下命令：

dlib_reproduce --model_num=<?> --study=fairness_study_v1

其中<?>应替换为0到12,599之间的模型索引， 对应于要训练的模型ID。

如果您只想使用论文中的评估协议重新评估已训练的模型，可以使用以下命令：

dlib_reproduce --model_dir=<model_output_directory> --output_directory=<output> --study=fairness_study_v1

UDR实验

该库还包含用于以下论文提出的无监督解纠缠排名（UDR）方法的代码，位于disentanglement_lib/bin/dlib_udr：

变分解纠缠表征学习的无监督模型选择 Sunny Duan、Loic Matthey、Andre Saraiva、Nicholas Watters、Christopher P. Burgess、Alexander Lerchner、Irina Higgins。

UDR可用于新训练的模型（例如通过运行dlib_reproduce获得的模型）或现有的预训练模型。模型训练完成后，可以通过以下命令计算其UDR分数：

dlib_udr --model_dirs=<model_output_directory1>,<model_output_directory2> \
  --output_directory=<output>

分数将导出到<output>/results/aggregate/evaluation.json，存储在model_scores属性下。分数将按照输入模型目录的顺序呈现。

弱监督实验

该库还包含用于以下论文提出的弱监督解纠缠方法的代码，位于disentanglement_lib/bin/dlib_reproduce_weakly_supervised：

无需妥协的弱监督解纠缠 Francesco Locatello、Ben Poole、Gunnar Rätsch、Bernhard Schölkopf、Olivier Bachem、Michael Tschannen。

dlib_reproduce_weakly_supervised --output_directory=<output> \
   --gin_model_config_dir=<dir> \
   --gin_model_config_name=<name> \
   --gin_postprocess_config_glob=<postprocess_configs> \
   --gin_evaluation_config_glob=<eval_configs> \
   --pipeline_seed=<seed>

半监督实验

该库还包含用于以下论文提出的半监督解纠缠方法的代码，位于disentanglement_lib/bin/dlib_reproduce_semi_supervised：

利用少量标签解纠缠变化因素 Francesco Locatello、Michael Tschannen、Stefan Bauer、Gunnar Rätsch、Bernhard Schölkopf、Olivier Bachem。

dlib_reproduce_weakly_supervised --output_directory=<output> \
   --gin_model_config_dir=<dir> \
   --gin_model_config_name=<name> \
   --gin_postprocess_config_glob=<postprocess_configs> \
   --gin_evaluation_config_glob=<eval_configs> \
   --gin_validation_config_glob=<val_configs> \
   --pipeline_seed=<seed> \
   --eval_seed=<seed> \
   --supervised_seed=<seed> \
   --num_labelled_samples=<num> \
   --train_percentage=0.9 \
   --labeller_fn="@perfect_labeller"

反馈

请将任何反馈发送至bachem@google.com和francesco.locatello@tuebingen.mpg.de。

引用

如果您使用disentanglement_lib，请考虑引用：

@inproceedings{locatello2019challenging,
  title={挑战无监督解纠缠表征学习中的常见假设},
  author={Locatello, Francesco and Bauer, Stefan and Lucic, Mario and Raetsch, Gunnar and Gelly, Sylvain and Sch{\"o}lkopf, Bernhard and Bachem, Olivier},
  booktitle={国际机器学习大会},
  pages={4114--4124},
  year={2019}
}

这并非谷歌官方支持的产品。

disentanglement_lib 中文快速上手指南

环境准备

• 系统要求：Linux 系统（推荐使用 Google Cloud Deep Learning VM 或 Ubuntu 20.04+）
• Python 版本：Python 3.7+
• 前置依赖：TensorFlow（GPU 或 CPU 版）、NumPy、SciPy、scikit-learn、TFHub、Gin-config

💡 国内用户建议：使用清华镜像加速 pip 安装
在 pip install 命令前添加 -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库：

git clone https://github.com/google-research/disentanglement_lib.git
cd disentanglement_lib

2. 安装依赖（推荐使用 GPU 版本）：

pip install .[tf_gpu] -i https://pypi.tuna.tsinghua.edu.cn/simple

若无 GPU，使用：pip install .[tf]

3. 验证安装：

dlib_tests

4. 下载数据集（自动下载 dSprites、Color-dSprites、Noisy-dSprites、SmallNORB、Cars3D）：

dlib_download_data

5. （可选）设置数据路径环境变量，避免每次手动指定：

export DISENTANGLEMENT_LIB_DATA=<你的数据目录路径>

将上述命令添加至 ~/.bashrc 以永久生效。

基本使用

最简示例：复现一个预设模型

运行以下命令，训练并评估第 0 号模型（约数小时）：

dlib_reproduce --model_num=0

训练完成后，结果将保存在 output/0/ 目录中，包含：

• 训练模型（TFHub 模块）
• 评估指标（JSON 格式）
• 可视化 GIF（在 visualizations/ 子目录）

查看模型配置（不训练）

如仅想查看配置而不训练：

dlib_reproduce --model_num=0 --only_print

聚合所有实验结果

训练多个模型后，聚合所有结果为一个 JSON 文件：

dlib_aggregate_results

生成的 results.json 可用 Pandas 加载分析：

from disentanglement_lib.utils import aggregate_results
df = aggregate_results.load_aggregated_json_results("results.json")
print(df.head())

✅ 推荐：首次使用建议从 --model_num=0 开始，快速验证环境是否正常。