文档

pythae

该库以统一的实现方式实现了几种最常见的（变分）自编码器模型。特别地，它提供了通过使用相同的自动编码神经网络架构来训练这些模型，从而进行基准实验和比较的可能性。其“自定义自编码器”功能允许您使用自己的数据以及自定义的编码器和解码器神经网络来训练这些模型。该库集成了诸如 wandb、mlflow 或 comet-ml 🧪 等实验监控工具，并支持在几行代码内从 HuggingFace Hub 🤗 上共享和加载模型。

新闻 📢

自 v0.1.0 版本起，Pythae 现已支持使用 PyTorch 的 DDP 进行分布式训练。现在您可以更快地在更大的数据集上训练您喜爱的 VAE，而且仍然只需几行代码。 请参阅我们的加速 基准测试。

快速访问：

• 安装
• 已实现的模型 / 已实现的采样器
• 可重复性声明 / 结果呈现方式
• 模型训练 / 数据生成 / 自定义网络架构 / 分布式训练
• 与 🤗 Hub 共享模型 / 使用 wandb 监控实验 / 使用 mlflow 监控实验 / 使用 comet_ml 监控实验
• 教程 / 文档
• 贡献 🚀 / 问题 🛠️
• 引用本仓库

安装

要安装该库的最新稳定版，请使用 pip 运行以下命令：

$ pip install pythae

要安装该库的最新 GitHub 版本，请使用 pip 运行以下命令：

$ pip install git+https://github.com/clementchadebec/benchmark_VAE.git

或者，您也可以克隆 GitHub 仓库以访问测试、教程和脚本。

$ git clone https://github.com/clementchadebec/benchmark_VAE.git

然后进入目录并安装库：

$ cd benchmark_VAE
$ pip install -e .

可用模型

以下是当前库中已实现的模型列表。

模型	训练示例	论文	官方实现
自编码器 (AE)
变分自编码器 (VAE)		链接
Beta 变分自编码器 (BetaVAE)		链接
变分自编码器结合线性归一化流 (VAE_LinNF)		链接
变分自编码器结合逆向自回归流 (VAE_IAF)		链接	链接
解耦合 Beta 变分自编码器 (DisentangledBetaVAE)		链接
因子分解解耦 (FactorVAE)		链接
Beta-TC-VAE (BetaTCVAE)		链接	链接
重要性加权自编码器 (IWAE)		链接	链接
多重重要性加权自编码器 (MIWAE)		链接
部分重要性加权自编码器 (PIWAE)		链接
组合重要性加权自编码器 (CIWAE)		链接
基于感知度量相似性 (MSSSIM) 的变分自编码器		链接
Wasserstein 自编码器 (WAE)		链接	链接
信息变分自编码器 (INFOVAE_MMD)		链接
VAMP 自编码器 (VAMP)		链接	链接
超球面变分自编码器 (SVAE)		链接	链接
波兰克圆盘变分自编码器 (PoincareVAE)		链接	链接
对抗自编码器 (Adversarial_AE)		链接
变分自编码器 GAN (VAEGAN) 🥗		链接	链接
向量量化变分自编码器 (VQVAE)		链接	链接
哈密顿变分自编码器 (HVAE)		链接	链接
使用 L2 解码器参数正则化的自编码器 (RAE_L2)		链接	链接
使用梯度惩罚正则化的自编码器 (RAE_GP)		链接	链接
黎曼哈密顿变分自编码器 (RHVAE)		链接	链接
层次残差量化 (HRQVAE)		链接	链接

请参阅重建和生成结果，了解所有上述模型的表现

可用的采样器

以下是当前库中已实现的模型列表。

采样器	模型	论文	官方实现
正态先验 (NormalSampler)	所有模型	链接
高斯混合 (GaussianMixtureSampler)	所有模型	链接	链接
两阶段VAE采样器 (TwoStageVAESampler)	所有基于VAE的模型	链接	链接
单位球面均匀采样器 (HypersphereUniformSampler)	SVAE	链接	链接
波兰卡圆盘采样器 (PoincareDiskSampler)	PoincareVAE	链接	链接
VAMP先验采样器 (VAMPSampler)	VAMP	链接	链接
流形采样器 (RHVAESampler)	RHVAE	链接	链接
掩码自回归流采样器 (MAFSampler)	所有模型	链接	链接
逆向自回归流采样器 (IAFSampler)	所有模型	链接	链接
PixelCNN (PixelCNNSampler)	VQVAE	链接

可复现性

我们通过复现原始论文中的一些结果来验证实现的正确性，前提是官方代码已发布，或者论文的实验部分提供了足够详细的信息。更多详情请参阅可复现性。

启动模型训练

要启动模型训练，只需调用一个 TrainingPipeline 实例即可。

>>> from pythae.pipelines import TrainingPipeline
>>> from pythae.models import VAE, VAEConfig
>>> from pythae.trainers import BaseTrainerConfig

>>> # 设置训练配置
>>> my_training_config = BaseTrainerConfig(
...	output_dir='my_model',
...	num_epochs=50,
...	learning_rate=1e-3,
...	per_device_train_batch_size=200,
...	per_device_eval_batch_size=200,
...	train_dataloader_num_workers=2,
...	eval_dataloader_num_workers=2,
...	steps_saving=20,
...	optimizer_cls="AdamW",
...	optimizer_params={"weight_decay": 0.05, "betas": (0.91, 0.995)},
...	scheduler_cls="ReduceLROnPlateau",
...	scheduler_params={"patience": 5, "factor": 0.5}
... )
>>> # 设置模型配置 
>>> my_vae_config = model_config = VAEConfig(
...	input_dim=(1, 28, 28),
...	latent_dim=10
... )
>>> # 构建模型
>>> my_vae_model = VAE(
...	model_config=my_vae_config
... )
>>> # 构建流水线
>>> pipeline = TrainingPipeline(
... 	training_config=my_training_config,
... 	model=my_vae_model
... )
>>> # 启动流水线
>>> pipeline(
...	train_data=your_train_data, # 必须是 torch.Tensor、np.array 或 torch 数据集
...	eval_data=your_eval_data # 必须是 torch.Tensor、np.array 或 torch 数据集
... )

训练结束后，最佳模型权重、模型配置和训练配置将存储在 my_model/MODEL_NAME_training_YYYY-MM-DD_hh-mm-ss 文件夹中的 final_model 目录下（其中 my_model 是 BaseTrainerConfig 的 output_dir 参数）。如果进一步设置了 steps_saving 参数，则还会在 my_model/MODEL_NAME_training_YYYY-MM-DD_hh-mm-ss 中出现名为 checkpoint_epoch_k 的文件夹，其中包含第 k 个 epoch 时的最佳模型权重、优化器、调度器、配置和训练配置。

在基准数据集上启动训练

我们还提供了一个训练脚本示例此处，可用于在基准数据集（mnist、cifar10、celeba 等）上训练模型。该脚本可以通过以下命令行启动：

python training.py --dataset mnist --model_name ae --model_config 'configs/ae_config.json' --training_config 'configs/base_training_config.json'

有关此脚本的更多详细信息，请参阅 README.md。

启动数据生成

使用 GenerationPipeline

从已训练模型中启动数据生成的最简单方法是使用 Pythae 内置的 GenerationPipeline。假设您想使用 MAFSampler 生成 100 个样本，您只需执行以下步骤：1) 重新加载已训练的模型，2) 定义采样器的配置，3) 创建并启动 GenerationPipeline，如下所示：

>>> from pythae.models import AutoModel
>>> from pythae.samplers import MAFSamplerConfig
>>> from pythae.pipelines import GenerationPipeline
>>> # 恢复已训练的模型
>>> my_trained_vae = AutoModel.load_from_folder(
...	'path/to/your/trained/model'
... )
>>> my_sampler_config = MAFSamplerConfig(
...	n_made_blocks=2,
...	n_hidden_in_made=3,
...	hidden_size=128
... )
>>> # 构建流水线
>>> pipe = GenerationPipeline(
...	model=my_trained_vae,
...	sampler_config=my_sampler_config
... )
>>> # 启动数据生成
>>> generated_samples = pipe(
...	num_samples=args.num_samples,
...	return_gen=True, # 如果为假则不返回任何内容
...	train_data=train_data, # 用于拟合采样器
...	eval_data=eval_data, # 用于拟合采样器
...	training_config=BaseTrainerConfig(num_epochs=200) # 用于拟合采样器的训练配置
... )

使用采样器

或者，你也可以直接通过采样器从训练好的模型中启动数据生成过程。例如，要使用你的采样器生成新数据，可以运行以下代码：

>>> from pythae.models import AutoModel
>>> from pythae.samplers import NormalSampler
>>> # 获取训练好的模型
>>> my_trained_vae = AutoModel.load_from_folder(
...	'path/to/your/trained/model'
... )
>>> # 定义你的采样器
>>> my_samper = NormalSampler(
...	model=my_trained_vae
... )
>>> # 生成样本
>>> gen_data = my_samper.sample(
...	num_samples=50,
...	batch_size=10,
...	output_dir=None,
...	return_gen=True
... )

如果你将 output_dir 设置为一个特定路径，生成的图像将会以 .png 文件的形式保存，文件名分别为 00000000.png, 00000001.png 等等。

只要模型适配，采样器就可以用于任何模型。例如，GaussianMixtureSampler 实例可以用于任何模型，而 VAMPSampler 只能与 VAMP 模型一起使用。请查看 此处 以了解哪些采样器适用于你的模型。请注意，某些采样器，比如 GaussianMixtureSampler，在使用前可能需要调用 fit 方法进行拟合。以下是 GaussianMixtureSampler 的示例：

>>> from pythae.models import AutoModel
>>> from pythae.samplers import GaussianMixtureSampler, GaussianMixtureSamplerConfig
>>> # 获取训练好的模型
>>> my_trained_vae = AutoModel.load_from_folder(
...	'path/to/your/trained/model'
... )
>>> # 定义你的采样器
... gmm_sampler_config = GaussianMixtureSamplerConfig(
...	n_components=10
... )
>>> my_samper = GaussianMixtureSampler(
...	sampler_config=gmm_sampler_config,
...	model=my_trained_vae
... )
>>> # 拟合采样器
>>> gmm_sampler.fit(train_dataset)
>>> # 生成样本
>>> gen_data = my_samper.sample(
...	num_samples=50,
...	batch_size=10,
...	output_dir=None,
...	return_gen=True
... )

自定义自编码器架构

Pythae 提供了在 VAE 模型中定义自定义神经网络的可能性。例如，假设你想训练一个带有特定编码器和解码器的 Wassertstein AE，你可以这样做：

>>> from pythae.models.nn import BaseEncoder, BaseDecoder
>>> from pythae.models.base.base_utils import ModelOutput
>>> class My_Encoder(BaseEncoder):
...	def __init__(self, args=None): # Args 是一个 ModelConfig 实例
...		BaseEncoder.__init__(self)
...		self.layers = my_nn_layers()
...		
...	def forward(self, x:torch.Tensor) -> ModelOutput:
...		out = self.layers(x)
...		output = ModelOutput(
...			embedding=out # 将编码器的输出放入 ModelOutput 实例中 
...		)
...		return output
...
... class My_Decoder(BaseDecoder):
...	def __init__(self, args=None):
...		BaseDecoder.__init__(self)
...		self.layers = my_nn_layers()
...		
...	def forward(self, x:torch.Tensor) -> ModelOutput:
...		out = self.layers(x)
...		output = ModelOutput(
...			reconstruction=out # 将解码器的输出放入 ModelOutput 实例中
...		)
...		return output
...
>>> my_encoder = My_Encoder()
>>> my_decoder = My_Decoder()

然后构建模型：

>>> from pythae.models import WAE_MMD, WAE_MMD_Config
>>> # 设置模型配置 
>>> my_wae_config = model_config = WAE_MMD_Config(
...	input_dim=(1, 28, 28),
...	latent_dim=10
... )
...
>>> # 构建模型
>>> my_wae_model = WAE_MMD(
...	model_config=my_wae_config,
...	encoder=my_encoder, # 在构建模型时传入你的编码器
...	decoder=my_decoder # 在构建模型时传入你的解码器
... )

重要提示 1：对于所有基于 AE 的模型（AE、WAE、RAE_L2、RAE_GP），编码器和解码器都必须返回一个 ModelOutput 实例。对于编码器，ModelOutput 实例必须在 embedding 键下包含嵌入向量。对于解码器，ModelOutput 实例必须在 reconstruction 键下包含重建结果。

重要提示 2：对于所有基于 VAE 的模型（VAE、BetaVAE、IWAE、HVAE、VAMP、RHVAE），编码器和解码器都必须返回一个 ModelOutput 实例。对于编码器，ModelOutput 实例必须分别在 embedding 和 log_covariance 键下包含嵌入向量和对数协方差矩阵（形状为 batch_size × latent_space_dim）。对于解码器，ModelOutput 实例必须在 reconstruction 键下包含重建结果。

使用基准神经网络

你还可以找到针对最常见数据集（如 MNIST、CIFAR、CELEBA 等）的预定义神经网络架构，可以通过以下方式加载：

>>> from pythae.models.nn.benchmark.mnist import (
...	Encoder_Conv_AE_MNIST, # 用于基于 AE 的模型（仅返回嵌入）
...	Encoder_Conv_VAE_MNIST, # 用于基于 VAE 的模型（返回嵌入和对数协方差）
...	Decoder_Conv_AE_MNIST
... )

将 mnist 替换为 cifar 或 celeba，即可访问其他神经网络。

使用 Pythae 进行分布式训练

自 v0.1.0 起，Pythae 现已支持使用 PyTorch 的 DDP 进行分布式训练。这使你能够利用多 GPU 和/或多节点训练，更快地在更大的数据集上训练你喜欢的 VAE。

为此，你可以编写一个 Python 脚本，然后由启动程序（如集群上的 srun）来运行该脚本。脚本中唯一需要做的就是在训练配置中直接指定与分布式环境相关的参数，如下所示：

>>> training_config = BaseTrainerConfig(
...     num_epochs=10,
...     learning_rate=1e-3,
...     per_device_train_batch_size=64,
...     per_device_eval_batch_size=64,
...     train_dataloader_num_workers=8,
...     eval_dataloader_num_workers=8,
...     dist_backend="nccl", # 分布式后端
...     world_size=8 # 使用的 GPU 数量（节点数 × 每个节点的 GPU 数量），
...     rank=5 # 全局 GPU ID，
...     local_rank=1 # 节点内的 GPU ID，
...     master_addr="localhost" # 主节点地址，
...     master_port="12345" # 主节点端口，
... )

请参阅此 示例脚本，其中定义了一个在 ImageNet 数据集上进行多 GPU VQVAE 训练的脚本。请注意，分布式环境变量（world_size、rank 等）的获取方式可能因你使用的集群和启动程序而异。

基准测试

以下是使用 Pythae 在 V100 16GB GPU 上对 MNIST 数据集训练 100 个 epoch 的向量量化变分自编码器 (VQ-VAE)、在 FFHQ（1024×1024 图像）上训练 50 个 epoch，以及在 V100 32GB GPU 上对 ImageNet-1k 训练 20 个 epoch 的训练时间。

	训练数据	1 张 GPU	4 张 GPU	2×4 张 GPU
MNIST (VQ-VAE)	28×28 图像（5 万张）	235.18 秒	62.00 秒	35.86 秒
FFHQ 1024×1024 (VQVAE)	1024×1024 RGB 图像（6 万张）	19 小时 1 分钟	5 小时 6 分钟	2 小时 37 分钟
ImageNet-1k 128×128 (VQVAE)	128×128 RGB 图像（约 120 万张）	6 小时 25 分钟	1 小时 41 分钟	51 分钟 26 秒

对于每个数据集，我们提供了基准测试脚本，可在这里找到。

使用 HuggingFace Hub 分享你的模型 🤗

Pythae 还允许你将模型分享到 HuggingFace Hub。为此你需要：

• 一个有效的 HuggingFace 账号
• 在你的虚拟环境中安装了 huggingface_hub 包。如果没有，可以通过以下命令安装：

$ python -m pip install huggingface_hub

• 使用以下命令登录你的 HuggingFace 账号：

$ huggingface-cli login

将模型上传到 Hub

任何 Pythae 模型都可以通过 push_to_hf_hub 方法轻松上传：

>>> my_vae_model.push_to_hf_hub(hf_hub_path="your_hf_username/your_hf_hub_repo")

注意： 如果 your_hf_hub_repo 已经存在且不为空，文件将会被覆盖。如果该仓库不存在，系统会创建一个同名的文件夹。

从 Hub 下载模型

同样地，你可以直接从 Hub 下载或重新加载任何 Pythae 模型，只需使用 load_from_hf_hub 方法：

>>> from pythae.models import AutoModel
>>> my_downloaded_vae = AutoModel.load_from_hf_hub(hf_hub_path="path_to_hf_repo")

使用 wandb 监控你的实验 🧪

Pythae 还集成了实验跟踪工具 wandb，允许用户存储配置、监控训练过程，并通过图形界面比较不同运行的结果。要使用此功能，你需要：

• 一个有效的 wandb 账号
• 在你的虚拟环境中安装了 wandb 包。如果没有，可以通过以下命令安装：

$ pip install wandb

• 使用以下命令登录你的 wandb 账号：

$ wandb login

创建 WandbCallback

在 Pythae 中使用 wandb 启动实验监控非常简单。用户只需创建一个 WandbCallback 实例……

>>> # 创建回调
>>> from pythae.trainers.training_callbacks import WandbCallback
>>> callbacks = [] # TrainingPipeline 需要一个回调列表
>>> wandb_cb = WandbCallback() # 构建回调
>>> # 设置回调
>>> wandb_cb.setup(
...	training_config=your_training_config, # 训练配置
...	model_config=your_model_config, # 模型配置
...	project_name="your_wandb_project", # 指定你的 wandb 项目
...	entity_name="your_wandb_entity", # 指定你的 wandb 实体
... )
>>> callbacks.append(wandb_cb) # 添加到回调列表

……然后将其传递给 TrainingPipeline。

>>> pipeline = TrainingPipeline(
...	training_config=config,
...	model=model
... )
>>> pipeline(
...	train_data=train_dataset,
...	eval_data=eval_dataset,
...	callbacks=callbacks # 将回调传递给 TrainingPipeline，大功告成！
... )
>>> # 你可以登录 https://wandb.ai/your_wandb_entity/your_wandb_project 来监控你的训练

请参阅详细教程

使用 mlflow 监控你的实验 🧪

Pythae 还集成了实验跟踪工具 mlflow，允许用户存储配置、监控训练过程，并通过图形界面比较不同运行的结果。要使用此功能，你需要：

• 在你的虚拟环境中安装了 mlflow 包。如果没有，可以通过以下命令安装：

$ pip install mlflow

创建 MLFlowCallback

在 Pythae 中使用 mlflow 启动实验监控非常简单。用户只需创建一个 MLFlowCallback 实例……

>>> # 创建回调
>>> from pythae.trainers.training_callbacks import MLFlowCallback
>>> callbacks = [] # TrainingPipeline 需要一个回调列表
>>> mlflow_cb = MLFlowCallback() # 构建回调
>>> # 设置回调
>>> mlflow_cb.setup(
...	training_config=your_training_config, # 训练配置
...	model_config=your_model_config, # 模型配置
...	run_name="mlflow_cb_example", # 指定你的 mlflow 运行名称
... )
>>> callbacks.append(mlflow_cb) # 添加到回调列表

……然后将其传递给 TrainingPipeline。

>>> pipeline = TrainingPipeline(
...	training_config=config，
...	model=model
... )
>>> pipeline(
...	train_data=train_dataset，
...	eval_data=eval_dataset，
...	callbacks=callbacks # 将回调传递给 TrainingPipeline，大功告成！
... )

你可以在包含 ./mlruns 的目录中运行以下命令来可视化指标：

$ mlflow ui

请参阅详细教程

使用 comet_ml 监控你的实验 🧪

Pythae 还集成了实验跟踪工具 comet_ml，允许用户存储配置、监控训练过程，并通过图形界面比较不同运行的结果。要使用此功能，你需要：

• 在你的虚拟环境中安装了 comet_ml 包。如果没有，可以通过以下命令安装：

$ pip install comet_ml

创建 CometCallback

在 Pythae 中使用 comet_ml 启动实验监控非常简单。用户只需创建一个 CometCallback 实例……

>>> # 创建回调
>>> from pythae.trainers.training_callbacks import CometCallback
>>> callbacks = [] # TrainingPipeline 需要一个回调列表
>>> comet_cb = CometCallback() # 构建回调
>>> # 设置回调
>>> comet_cb.setup(
...	training_config=training_config, # 训练配置
...	model_config=model_config, # 模型配置
...	api_key="your_comet_api_key", # 指定你的 comet API 密钥
...	project_name="your_comet_project", # 指定你的 comet 项目
...	#offline_run=True, # 以离线模式运行
...	#offline_directory='my_offline_runs' # 设置用于存储离线运行的目录
... )
>>> callbacks.append(comet_cb) # 添加到回调列表

……然后将其传递给 TrainingPipeline。

>>> pipeline = TrainingPipeline(
...	training_config=config，
...	model=model
... )
>>> pipeline(
...	train_data=train_dataset，
...	eval_data=eval_dataset，
...,callbacks=callbacks # 将回调传递给 TrainingPipeline，大功告成！
... )
>>> # 你可以登录 https://comet.com/your_comet_username/your_comet_project 来监控你的训练

请参阅详细教程

获取代码

为了帮助您理解 pythae 的工作原理以及如何使用该库训练您的模型，我们还提供了以下教程：

• making_your_own_autoencoder.ipynb 向您展示如何将您自己的网络传递给 pythae 中实现的模型

• custom_dataset.ipynb 向您展示如何在 pythae 中实现的任何模型中使用自定义数据集

• hf_hub_models_sharing.ipynb 向您展示如何上传和下载 HuggingFace Hub 上的模型

• wandb_experiment_monitoring.ipynb 向您展示如何使用 wandb 监控实验

• mlflow_experiment_monitoring.ipynb 向您展示如何使用 mlflow 监控实验

• comet_experiment_monitoring.ipynb 向您展示如何使用 comet_ml 监控实验

• models_training 文件夹提供了笔记本，展示了如何训练每个已实现的模型，以及如何使用 pythae.samplers 从中采样。

• scripts 文件夹特别提供了一个训练脚本示例，用于在基准数据集（mnist、cifar10、celeba 等）上训练模型。

处理问题 🛠️

如果您在运行代码时遇到任何问题，或希望添加新的功能/模型，请在 github 上提交一个问题。

贡献 🚀

您想通过添加一个模型、采样器，或者只是修复一个 bug 来为这个库做出贡献吗？太棒了！非常感谢！请参阅 CONTRIBUTING.md，以了解主要的贡献指南。

结果

重建效果

首先让我们来看看从评估集抽取的重建样本。

模型	MNIST	CELEBA
评估数据
自编码器
变分自编码器
Beta-变分自编码器
线性流变分自编码器
IAF变分自编码器
解耦合Beta-变分自编码器
FactorVAE
BetaTCVAE
IWAE
MSSSIM_VAE
WAE
INFO VAE
VAMP
SVAE
对抗自编码器
VAE_GAN
VQVAE
HVAE
RAE_L2
RAE_GP
黎曼哈密顿变分自编码器 (RHVAE)

---

生成效果

在这里，我们展示了使用库中实现的各个模型以及不同采样器生成的样本。

模型	MNIST	CELEBA
AE + 高斯混合采样器
VAE + 正态采样器
VAE + 高斯混合采样器
VAE + 两阶段VAE采样器
VAE + MAF采样器
Beta-VAE + 正态采样器
VAE Lin NF + 正态采样器
VAE IAF + 正态采样器
解耦合Beta-VAE + 正态采样器
FactorVAE + 正态采样器
BetaTCVAE + 正态采样器
IWAE + 正态采样器
MSSSIM_VAE + 正态采样器
WAE + 正态采样器
INFO VAE + 正态采样器
SVAE + 超球面均匀采样器
VAMP + VAMP采样器
对抗自编码器 + 正态采样器
VAEGAN + 正态采样器
VQVAE + MAF采样器
HVAE + 正态采样器
RAE_L2 + 高斯混合采样器
RAE_GP + 高斯混合采样器
黎曼哈密顿VAE (RHVAE) + RHVAE采样器

引用

如果您觉得这项工作有用，或在您的研究中使用了它，请考虑引用我们。

@inproceedings{chadebec2022pythae,
 author = {Chadebec, Cl\'{e}ment and Vincent, Louis and Allassonniere, Stephanie},
 booktitle = {神经信息处理系统进展},
 editor = {S. Koyejo and S. Mohamed and A. Agarwal and D. Belgrave and K. Cho and A. Oh},
 pages = {21575--21589},
 publisher = {柯兰联合公司},
 title = {Pythae：在 Python 中统一生成式自编码器——一个基准测试用例},
 volume = {35},
 year = {2022}
}

Pythae (benchmark_VAE) 快速上手指南

Pythae 是一个统一的变分自编码器（VAE）库，集成了多种主流 VAE 模型。它支持在相同的网络架构下进行基准测试和对比实验，允许用户自定义编码器和解码器，并原生支持分布式训练、实验监控（WandB, MLflow, Comet）以及 HuggingFace Hub 模型共享。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：3.7, 3.8, 3.9 或更高版本
• 核心依赖：PyTorch (库会自动处理相关依赖)
• 可选依赖：
  • 分布式训练：需配置支持 NCCL 的 GPU 环境
  • 实验监控：wandb, mlflow, 或 comet_ml

安装步骤

方式一：通过 Pip 安装（推荐）

安装最新稳定版：

pip install pythae

若需使用国内镜像加速安装：

pip install pythae -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从 GitHub 安装（获取最新特性）

如需体验最新功能（如最新的分布式训练优化），可直接从源码安装：

pip install git+https://github.com/clementchadebec/benchmark_VAE.git

或者克隆仓库以便运行示例和测试：

git clone https://github.com/clementchadebec/benchmark_VAE.git
cd benchmark_VAE
pip install -e .

基本使用

Pythae 的核心优势在于统一的接口。以下是最简单的训练流程示例，以标准的 VAE 模型为例。

1. 导入模块与配置

首先导入所需的模型类和训练配置类，并设置参数。

from pythae.models import VAE, VAEConfig
from pythae.trainers import BaseTrainer, BaseTrainerConfig

# 定义模型配置
model_config = VAEConfig(
    input_dim=(3, 28, 28),  # 输入数据维度 (Channels, Height, Width)
    latent_dim=10           # 潜在空间维度
)

# 初始化模型
model = VAE(model_config)

# 定义训练器配置
trainer_config = BaseTrainerConfig(
    output_dir="my_vae_experiment",
    train_batch_size=64,
    eval_batch_size=64,
    learning_rate=1e-3,
    num_epochs=10
)

2. 准备数据

Pythae 兼容标准的 PyTorch Dataset 和 DataLoader。

from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

# 加载数据集 (以 MNIST 为例)
train_dataset = MNIST(root="./data", train=True, download=True, transform=ToTensor())
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

eval_dataset = MNIST(root="./data", train=False, download=True, transform=ToTensor())
eval_loader = DataLoader(eval_dataset, batch_size=64, shuffle=False)

3. 启动训练

创建训练器并开始训练。

# 初始化训练器
trainer = BaseTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    config=trainer_config
)

# 开始训练
training_output = trainer.train()

4. 生成数据与保存

训练完成后，可以使用模型生成样本或保存到 HuggingFace Hub。

# 生成样本
generated_data = model.generate(num_samples=10)

# 保存到本地
model.save("my_saved_vae")

# 推送到 HuggingFace Hub (需先登录 huggingface-cli login)
# model.push_to_hub("your_username/my-vae-model")

提示：库中已内置了包括 BetaVAE, IWAE, VQVAE, WAE 等在内的 20+ 种模型，只需将上述代码中的 VAE 和 VAEConfig 替换为对应的模型类（如 BetaVAE, VQVAEConfig 等）即可无缝切换。