Pythae:一个统一的Python生成式自编码器库

benchmark_VAE

Pythae简介

Pythae是一个功能强大的Python库,旨在为各种生成式自编码器模型提供统一的实现和使用框架。它的出现解决了现有生成式自编码器实现分散、使用不便的问题,为研究人员和开发者提供了一个便捷的工具。

Pythae的主要特性

统一的实现: Pythae提供了多种常见生成式自编码器模型的统一实现,包括VAE、Beta-VAE、IWAE等。
易于使用的框架: 通过简单的API,用户可以轻松训练模型、生成数据。
可复现性: Pythae注重实验的可复现性,提供了详细的配置选项。
灵活性: 支持自定义网络架构,可以灵活应对不同需求。
分布式训练: 支持使用PyTorch的DDP进行分布式训练,提高训练效率。
实验监控: 集成了wandb、mlflow等实验监控工具。
模型共享: 支持通过HuggingFace Hub轻松共享和加载模型。

Pythae的主要组件

1. 模型库

Pythae实现了大量常见的生成式自编码器模型,主要包括:

自编码器(AE)
变分自编码器(VAE)
Beta变分自编码器(Beta-VAE)
具有线性正则化流的VAE(VAE_LinNF)
具有逆自回归流的VAE(VAE_IAF)
解缠Beta变分自编码器(DisentangledBetaVAE)
通过因子分解解缠(FactorVAE)
Beta-TC-VAE
重要性加权自编码器(IWAE)
多重重要性加权自编码器(MIWAE)
部分重要性加权自编码器(PIWAE)
组合重要性加权自编码器(CIWAE)
具有感知度量相似性的VAE(MSSSIM_VAE)
Wasserstein自编码器(WAE)
信息变分自编码器(INFOVAE_MMD)
VAMP自编码器(VAMP)
超球面VAE(SVAE)
Poincaré圆盘VAE(PoincareVAE)
对抗自编码器(Adversarial_AE)
变分自编码器GAN(VAEGAN)
矢量量化VAE(VQVAE)
哈密顿VAE(HVAE)
具有L2解码器参数的正则化AE(RAE_L2)
具有梯度惩罚的正则化AE(RAE_GP)
黎曼哈密顿VAE(RHVAE)
分层残差量化(HRQVAE)

这些模型涵盖了生成式自编码器领域的大部分主流算法,为用户提供了丰富的选择。

2. 采样器

除了模型库,Pythae还提供了多种采样器,用于从训练好的模型中生成新数据:

正态先验采样器(NormalSampler)
高斯混合采样器(GaussianMixtureSampler)
两阶段VAE采样器(TwoStageVAESampler)
单位球面均匀采样器(HypersphereUniformSampler)
Poincaré圆盘采样器(PoincareDiskSampler)
VAMP先验采样器(VAMPSampler)
流形采样器(RHVAESampler)
掩码自回归流采样器(MAFSampler)
逆自回归流采样器(IAFSampler)
PixelCNN采样器(PixelCNNSampler)

这些采样器可以与不同的模型配合使用,为数据生成提供了灵活的选择。

使用Pythae

安装

Pythae可以通过pip轻松安装:

pip install pythae

如果想使用最新的开发版本,可以直接从GitHub安装:

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

模型训练

使用Pythae训练模型非常简单,只需要几个步骤:

导入必要的模块
设置训练配置
设置模型配置
构建模型
创建训练管道
启动训练

下面是一个使用VAE模型的示例代码:

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 = 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 datasets
    eval_data=your_eval_data # 必须是torch.Tensor, np.array或torch datasets
)

训练完成后,最佳模型权重、模型配置和训练配置将存储在my_model/MODEL_NAME_training_YYYY-MM-DD_hh-mm-ss/final_model文件夹中。

数据生成

Pythae提供了两种方式来生成新数据:使用GenerationPipeline或直接使用采样器。

使用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=100,
    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
)