无条件图像生成

无条件图像生成模型在训练过程中不以文本或图像作为条件。它仅生成与训练数据分布相似的图像。

本指南将深入讲解 train_unconditional.py 训练脚本，帮助你熟悉该脚本，并了解如何将其适配到自身的使用场景中。

在运行脚本之前，请确保从源代码安装该库：

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

然后导航到包含训练脚本的示例文件夹，并安装所需的依赖项：

cd examples/unconditional_image_generation
pip install -r requirements.txt

🤗 Accelerate 是一个用于在多个 GPU/TPU 上训练或使用混合精度训练的库。它会根据你的硬件和环境自动配置训练设置。请查看 🤗 Accelerate 的快速入门以了解更多信息。

初始化一个 🤗 Accelerate 环境：

accelerate config

要设置一个默认的 🤗 Accelerate 环境，而不选择任何配置：

accelerate config default

或者，如果你的环境不支持交互式 shell（例如 notebook），你可以使用：

from accelerate.utils import write_basic_config

write_basic_config()

最后，如果您想在自己的数据集上训练模型，请查看创建训练数据集指南，了解如何创建可与训练脚本配合使用的数据集。

脚本参数

以下各部分重点介绍了训练脚本中对于理解如何修改脚本至关重要的内容，但并未详细涵盖脚本的每个方面。如果您希望了解更多，请随时阅读脚本，并告诉我们您是否有任何疑问或顾虑。

训练脚本提供了许多参数，帮助你自定义训练过程。所有参数及其说明都可以在 parse_args() 函数中找到。它为每个参数提供了默认值，例如训练批次大小和学习率，但你也可以在训练命令中设置自己的值。

例如，若要使用 bf16 格式通过混合精度加速训练，请在训练命令中添加 --mixed_precision 参数：

accelerate launch train_unconditional.py \
  --mixed_precision="bf16"

一些需要指定的基本且重要的参数包括：

--dataset_name：用于训练的 Hub 上的数据集名称或本地数据集路径
--output_dir：保存训练后模型的位置
--push_to_hub：是否将训练好的模型推送到 Hub
--checkpointing_steps：模型训练过程中保存检查点的频率；如果训练被中断，这非常有用，您可以在训练命令中添加 --resume_from_checkpoint 从该检查点继续训练

只需提供你的数据集，其余一切交由训练脚本处理！

训练脚本

用于预处理数据集和训练循环的代码位于 main() 函数中。如果需要调整训练脚本，请在此处进行修改。

如果您未提供模型配置，train_unconditional 脚本会初始化一个 UNet2DModel。您可以在此处配置 UNet：

model = UNet2DModel(
    sample_size=args.resolution,
    in_channels=3,
    out_channels=3,
    layers_per_block=2,
    block_out_channels=(128, 128, 256, 256, 512, 512),
    down_block_types=(
        "DownBlock2D",
        "DownBlock2D",
        "DownBlock2D",
        "DownBlock2D",
        "AttnDownBlock2D",
        "DownBlock2D",
    ),
    up_block_types=(
        "UpBlock2D",
        "AttnUpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
    ),
)

接下来，脚本初始化一个 scheduler 和一个 optimizer：

# Initialize the scheduler
accepts_prediction_type = "prediction_type" in set(inspect.signature(DDPMScheduler.__init__).parameters.keys())
if accepts_prediction_type:
    noise_scheduler = DDPMScheduler(
        num_train_timesteps=args.ddpm_num_steps,
        beta_schedule=args.ddpm_beta_schedule,
        prediction_type=args.prediction_type,
    )
else:
    noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule)

# Initialize the optimizer
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

然后它会加载数据集，你可以指定如何对其进行预处理：

dataset = load_dataset("imagefolder", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split="train")

augmentations = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
        transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)

最后，训练循环负责处理其余所有工作，例如为图像添加噪声、预测噪声残差、计算损失、在指定步数保存检查点，以及将模型保存并推送到 Hub。如果你想进一步了解训练循环的工作原理，请查看理解流水线、模型和调度器教程，该教程详细拆解了去噪过程的基本模式。

运行脚本

完成所有更改或对默认配置满意后，即可启动训练脚本！🚀

完整的训练过程在 4 块 V100 GPU 上需要 2 小时。

单 GPU

多 GPU

训练脚本会在您的仓库中创建并保存一个检查点文件。现在，您可以加载并使用训练好的模型进行推理：

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("anton-l/ddpm-butterflies-128").to("cuda")  # or "mps", "xpu", "cpu"
image = pipeline().images[0]

Update on GitHub

Diffusers

无条件图像生成

脚本参数

训练脚本

运行脚本