扩散模型实战（八）：微调扩散模型

  微调在LLM中并不是新鲜的概念，从头开始训练一个扩散模型需要很长的时间，特别是使用高分辨率图像训练。那么其实我们可以在已经训练好的”去噪“扩散模型基础上使用微调数据集进行二次微调训练。

       本文将介绍基于蝴蝶数据集上微调人脸生成的扩散模型：

一、环境准备

1.1 安装相关库

!pip install -qq diffusers datasets accelerate wandb open-clip-torch

1.2 登录Huggingface Hub

如果需要开源微调好的模型到Huggingface Hub上，那么需要使用如下代码登录，否则可忽略此步骤：

from huggingface_hub import notebook_login

notebook_login()

1.3 导入相关库

import numpy as np

import torch

import torch.nn.functional as F

import torchvision

from datasets import load_dataset

from diffusers import DDIMScheduler, DDPMPipeline

from matplotlib import pyplot as plt

from PIL import Image

from torchvision import transforms

from tqdm.auto import tqdm

device = (

    "mps"

    if torch.backends.mps.is_available()

    else "cuda"

    if torch.cuda.is_available()

    else "cpu"

)

二、导入预训练的扩散模型

下面我们导入人脸生成的扩散模型，观察一下生成的效果，代码如下：

image_pipe = DDPMPipeline.from_pretrained("google/ddpm-celebahq-256")

image_pipe.to(device);

查看生成的图像，代码如下：

images = image_pipe().images

images[0]

生成的效果虽然不错，但是速度稍微有点慢，其实有更快的采样器可以加速这一过程，比如下面介绍的DDIM

三、DDIM-更快的采样器

       在生成图像的每一步中，模型都会接收一个带有噪声的输入，并且需要预测这个噪声，以此来估计没有噪声的完整图像是什么。这个过程被称为采样过程，在Diffusers库中，采样通过调度器控制的，之前的文章中介绍过DDPMScheduler调度器，本文介绍的DDIMScheduler可以通过更少的迭代周期来产生很好的采样样本（1000多步采样不是必须的）。

# 创建一个新的调度器并设置推理迭代次数

scheduler = DDIMScheduler.from_pretrained("google/ddpm-celebahq-256")

scheduler.set_timesteps(num_inference_steps=40)

scheduler.timesteps

# 输出

tensor([975, 950, 925, 900, 875, 850, 825, 800, 775, 750, 725, 

        700, 675, 650, 625, 600, 575, 550, 525, 500, 475, 450, 425, 

        400, 375, 350, 325, 300, 275, 250, 225, 200, 175, 150, 

        125, 100,  75,  50,  25, 0])

    下面使用4幅随机噪声图像进行循环采样，并观察每一步的输入与输出的”去噪“图像，代码如下：

# 从随机噪声开始

x = torch.randn(4, 3, 256, 256).to(device) 

# batch size为4，三通道，长、宽均为256像素的一组图像

# 循环一整套时间步

for i, t in tqdm(enumerate(scheduler.timesteps)):

　

    # 准备模型输入：给“带躁”图像加上时间步信息

    model_input = scheduler.scale_model_input(x, t)

　

    # 预测噪声

    with torch.no_grad():

        noise_pred = image_pipe.unet(model_input, t)["sample"]

　

    # 使用调度器计算更新后的样本应该是什么样子

    scheduler_output = scheduler.step(noise_pred, t, x)

　

    # 更新输入图像

    x = scheduler_output.prev_sample

　

    # 时不时看一下输入图像和预测的“去噪”图像

    if i % 10 == 0 or i == len(scheduler.timesteps) - 1:

        fig, axs = plt.subplots(1, 2, figsize=(12, 5))

　

        grid = torchvision.utils.make_grid(x, nrow=4).permute(1, 2, 0)

        axs[0].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)

        axs[0].set_title(f"Current x (step {i})")

　

        pred_x0 = (

            scheduler_output.pred_original_sample

        ) 

        grid = torchvision.utils.make_grid(pred_x0, nrow=4).

           permute(1, 2, 0)

        axs[1].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)

        axs[1].set_title(f"Predicted denoised images (step {i})")

        plt.show()

  第二步生成图像的采样器是DDPMScheduler，我们可以使用新的DDIMScheduler来代替DDPMScheduler看看image_pipe生成的效果是否有提升，代码如下：

image_pipe.scheduler = scheduler

images = image_pipe(num_inference_steps=40).images

images[0]

  上述介绍了生成人脸的扩散模型以及生成的效果，也介绍了更快的采样器DDIMScheduler，下面我们使用蝴蝶数据集来微调人脸生成扩散模型：

四、微调人脸生成扩散模型

4.1 加载蝴蝶数据集

dataset_name = "huggan/smithsonian_butterflies_subset"

dataset = load_dataset(dataset_name, split="train")

image_size = 256

batch_size = 4

preprocess = transforms.Compose(

    [

        transforms.Resize((image_size, image_size)),

        transforms.RandomHorizontalFlip(),

        transforms.ToTensor(),

        transforms.Normalize([0.5], [0.5]),

    ]

)

　

def transform(examples):

    images = [preprocess(image.convert("RGB")) for image in 

        examples["image"]]

    return {"images": images}

　

dataset.set_transform(transform)

　

train_dataloader = torch.utils.data.DataLoader(

    dataset, batch_size=batch_size, shuffle=True

)

输出4幅蝴蝶图像，便于观察

print("Previewing batch:")

batch = next(iter(train_dataloader))

grid = torchvision.utils.make_grid(batch["images"], nrow=4)

plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5)

4.2 微调人脸生成扩散模型

num_epochs = 2

lr = 1e-5

grad_accumulation_steps = 2

　

optimizer = torch.optim.AdamW(image_pipe.unet.parameters(), lr=lr)

　

losses = []

　

for epoch in range(num_epochs):

    for step, batch in tqdm(enumerate(train_dataloader), 

      total=len(train_dataloader)):

        clean_images = batch["images"].to(device)

        # 随机生成一个噪声，稍后加到图像上

        noise = torch.randn(clean_images.shape).to(clean_images.

           device)

        bs = clean_images.shape[0]

　

        # 随机选取一个时间步

        timesteps = torch.randint(

            0,

            image_pipe.scheduler.num_train_timesteps,

            (bs,),

            device=clean_images.device,

        ).long()

　

        # 根据选中的时间步和确定的幅值，在干净图像上添加噪声

        # 此处为前向扩散过程

        noisy_images = image_pipe.scheduler.add_noise(clean_images, 

           noise, timesteps)

　

        # 使用“带噪”图像进行网络预测

        noise_pred = image_pipe.unet(noisy_images, timesteps, 

           return_dict=False)[0]

　

        # 对真正的噪声和预测的结果进行比较，注意这里是预测噪声

        loss = F.mse_loss(

            noise_pred, noise

        )

　

        # 保存损失值

        losses.append(loss.item())

　

        # 根据损失值更新梯度

        loss.backward()

　

        # 进行梯度累积，在累积到一定步数后更新模型参数

        if (step + 1) % grad_accumulation_steps == 0:

            optimizer.step()

            optimizer.zero_grad()

　

    print(

        f"Epoch {epoch} average loss: {sum(losses[-len(train_

           dataloader):])/len(train_dataloader)}"

    )

　

# 画出损失曲线，效果如图所示

plt.plot(losses) 

4.3 使用微调好的模型生成图像

x = torch.randn(8, 3, 256, 256).to(device)

for i, t in tqdm(enumerate(scheduler.timesteps)):

    model_input = scheduler.scale_model_input(x, t)

    with torch.no_grad():

        noise_pred = image_pipe.unet(model_input, t)["sample"]

    x = scheduler.step(noise_pred, t, x).prev_sample

grid = torchvision.utils.make_grid(x, nrow=4)

plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5)

从图中可以看出生成的图像有蝴蝶数据的风格。

4.4 保持微调好的扩散模型，并且上传到Huggingface Hub中

image_pipe.save_pretrained("my-finetuned-model")

from huggingface_hub import HfApi, ModelCard, create_repo, get_

   full_repo_name

# 配置Hugging Face Hub，上传文件

model_name = "ddpm-celebahq-finetuned-butterflies-2epochs" 

 # 使用@param 脚本程序对上传到

 # Hugging Face Hub的文件进行命名

local_folder_name = "my-finetuned-model" # @param脚本程序生成的名字，

# 你也可以通过 image_pipe.save_pretrained('savename')自行指定

description = "Describe your model here" # @param

hub_model_id = get_full_repo_name(model_name)

create_repo(hub_model_id)

api = HfApi()

api.upload_folder(

     folder_path=f"{local_folder_name}/scheduler",path_in_repo="", 

                    repo_id=hub_model_id )

api.upload_folder(

     folder_path=f"{local_folder_name}/unet", path_in_repo="", 

     repo_id=hub_model_id )

api.upload_file(

     path_or_fileobj=f"{local_folder_name}/model_index.json",

     path_in_repo="model_index.json",

     repo_id=hub_model_id,

)

　

# 添加一个模型卡片，这一步虽然不是必需的，但可以给他人提供一些模型描述信息

　

content = f"""

---

license: mit

tags:

- pytorch

- diffusers

- unconditional-image-generation

- diffusion-models-class

---

# 用法

from diffusers import DDPMPipeline

pipeline = DDPMPipeline.from_pretrained(' {hub_model_id}')

image = pipeline().images[0]

image

'''

"""

card = ModelCard(content)

card.push_to_hub(hub_model_id)

微调Trick：

• 设置合适的batch_size，batch_size要在不超过GPU显存的前提下，尽量大一些，这样可以提高GPU计算效果；如果特别小，可以采用梯度累积的方式来更新模型参数，达到和大batch_size类似的效果，也就是多运行几次loss.backward()，再调用optimizer.step()和optimizer.zero_grad()；
• 训练过程中，要时不时生成一些图像样本来观察模型性能；
• 训练过程中，可以把损失值、生成的图像样本等信息记录在日志中，可以使用Weights and Biases、TensorBoard等工具；

文章转自微信公众号@ArronAI