使用HuggingFace Diffusers训练无条件图像生成模型

时飞城Herdsman

于 2025-06-01 09:02:36 发布

阅读量249

点赞数 4

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00909/article/details/148360536

使用HuggingFace Diffusers训练无条件图像生成模型

diffusers Diffusers：在PyTorch中用于图像和音频生成的最先进扩散模型。项目地址: https://gitcode.com/gh_mirrors/di/diffusers

Diffusers是HuggingFace推出的一个专注于扩散模型的Python库，它提供了训练和使用各种扩散模型的工具。本文将详细介绍如何使用Diffusers库训练无条件图像生成模型。

环境准备

在开始训练前，需要确保已安装必要的依赖项。建议创建一个新的Python虚拟环境以避免依赖冲突。

安装Diffusers库

推荐从源码安装最新版本的Diffusers库，这样可以确保获得最新的功能和修复：

pip install diffusers[training]

安装示例特定依赖

进入示例目录后，安装训练脚本所需的额外依赖：

pip install -r requirements.txt

配置Accelerate环境

Diffusers使用Accelerate库来简化多GPU训练配置。运行以下命令初始化Accelerate环境：

accelerate config

按照提示选择适合你硬件配置的选项。

训练无条件图像生成模型

无条件图像生成是指模型在没有任何条件输入的情况下生成图像。下面介绍两种常见数据集的训练方法。

在Oxford Flowers数据集上训练

Oxford Flowers是一个包含102种花卉类别的图像数据集。以下是训练命令示例：

accelerate launch train_unconditional.py \
  --dataset_name="flowers-102-categories" \
  --resolution=64 \
  --center_crop \
  --random_flip \
  --output_dir="ddpm-ema-flowers-64" \
  --train_batch_size=16 \
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
  --use_ema \
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision=no

关键参数说明：

resolution: 输入图像的分辨率
center_crop: 对图像进行中心裁剪
random_flip: 随机水平翻转图像作为数据增强
use_ema: 使用指数移动平均来稳定训练
lr_warmup_steps: 学习率预热步数

在4块V100 GPU上，完整训练大约需要2小时。

在Pokemon数据集上训练

Pokemon数据集包含各种宝可梦角色的图像。训练命令类似：

accelerate launch train_unconditional.py \
  --dataset_name="pokemon" \
  --resolution=64 \
  --center_crop \
  --random_flip \
  --output_dir="ddpm-ema-pokemon-64" \
  --train_batch_size=16 \
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
  --use_ema \
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision=no

多GPU训练

Accelerate库简化了多GPU训练配置。以下是使用混合精度训练和多GPU的示例：

accelerate launch --mixed_precision="fp16" --multi_gpu train_unconditional.py \
  --dataset_name="pokemon" \
  --resolution=64 \
  --center_crop \
  --random_flip \
  --output_dir="ddpm-ema-pokemon-64" \
  --train_batch_size=16 \
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
  --use_ema \
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision="fp16" \
  --logger="wandb"

要使用Weights & Biases记录训练过程，需要先安装wandb库：

pip install wandb

使用自定义数据集

有两种方式使用自定义数据集：

1. 提供本地文件夹

确保图像文件按以下结构组织：

data_dir/
  ├── image1.png
  ├── image2.jpg
  └── ...

然后运行：

accelerate launch train_unconditional.py \
    --train_data_dir <path-to-train-directory> \
    <other-arguments>

2. 上传数据集到模型中心

可以使用Datasets库的ImageFolder功能轻松上传数据集：

from datasets import load_dataset

# 从本地文件夹创建数据集
dataset = load_dataset("imagefolder", data_dir="path_to_your_folder")

# 上传到模型中心
dataset.push_to_hub("your_dataset_name")

上传后，只需在训练时指定数据集名称即可。