01_finetuning_and_guidance_CN

# @markdown load and prepare a dataset:
# Not on Colab? Comments with #@ enable UI tweaks like headings or user inputs
# but can safely be ignored if you're working on a different platform.

dataset_name = "fashion_mnist"  # @param
dataset = load_dataset(dataset_name, split="train")
image_size = 128  # @param
batch_size = 4  # @param
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
)

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);
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

Downloading builder script:   0%|          | 0.00/4.83k [00:00

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11
• 12
• 13
• 14
• 15
• 16
• 17
• 18
• 19
• 20
• 21
• 22
• 23
• 24
• 25
• 26
• 27
• 28
• 29
• 30
• 31
• 32
• 33
• 34
• 35
• 36
• 37
• 38
• 39
• 40
• 41
• 42
• 43
• 44

考虑因素1： 我们这里使用的 batch size 很小（只有 4），因为我们的训练是基于较大的图片尺寸的（256px），并且我们的模型也很大，如果我们的 batch size 太高，GPU 的内存可能会不够用了。你可以减小图片尺寸，换取更大的 batch size 来加速训练，但这里的模型一开始都是基于生成 256px 尺寸的图片来设计和训练的。

现在我们看看训练循环。我们把要优化的目标参数设定为image_pipe.unet.parameters()，以此来更新预训练过的模型的权重。其它部分的代码基本上和第一单元例子中的对应部分一样。在 Colab 上跑的话，大约需要10分钟，你可以趁这个时间喝杯茶休息一下。

num_epochs = 2  # @param
lr = 1e-5  # 2param
grad_accumulation_steps = 2  # @param

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)
        # Sample noise to add to the images
        noise = torch.randn(clean_images.shape).to(clean_images.device)
        bs = clean_images.shape[0]

        # Sample a random timestep for each image
        timesteps = torch.randint(
            0,
            image_pipe.scheduler.num_train_timesteps,
            (bs,),
            device=clean_images.device,
        ).long()

        # Add noise to the clean images according to the noise magnitude at each timestep
        # (this is the forward diffusion process)
        noisy_images = image_pipe.scheduler.add_noise(clean_images, noise, timesteps)

        # Get the model prediction for the noise
        noise_pred = image_pipe.unet(noisy_images, timesteps, return_dict=False)[0]

        # Compare the prediction with the actual noise:
        loss = F.mse_loss(
            noise_pred, noise
        )  # NB - trying to predict noise (eps) not (noisy_ims-clean_ims) or just (clean_ims)

        # Store for later plotting
        losses.append(loss.item())

        # Update the model parameters with the optimizer based on this loss
        loss.backward(loss)

        # Gradient accumulation:
        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)}"
    )

# Plot the loss curve:
plt.plot(losses)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

  0%|          | 0/15000 [00:00]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

考虑因素2： 我们的损失值曲线简直像噪声一样混乱！这是因为每一次迭代我们都只用了四个训练样本，而且加到它们上面的噪声水平还都是随机挑选的。这对于训练来讲并不理想。一种弥补的措施是，我们使用一个非常小的学习率，限制每次更新的幅度。但我们还有一个更好的方法，既能得到和使用更大的 batch size 一样的收益，又不需要让我们的内存爆掉。

点击这里看看：gradient accumulation。如果我们多运行几次loss.backward()后再调用optimizer.step()和optimizer.zero_grad()，PyTorch 就会把梯度累积（加和）起来，这样多个批次的数据产生的更新信号就会被高效地融合在一起，产出一个单独的（更好的）梯度估计用于参数更新。这样做会减少参数更新的总次数，就正如我们使用更大的 batch size 时希望看到的一样。梯度累积是一个很多框架都会替你做的事情（比如这里：🤗 Accelerate makes this easy），但这里我们从头实现一遍也挺好的，因为这对你在 GPU 内存受限时训练模型非常有帮助。正如你在上面代码中看到的那样（在注释 # Gradient accumulation 后），其实也不需要你写很多代码。

# 练习：试试你能不能把梯度累积加到第一单元的训练循环中
# 怎么做呢？你应该怎么基于梯度累积的步数来调整学习率？
# 学习率应该和之前一样吗？
1
2
3

考虑因素3： 即使这样，我们的训练还是挺慢的，而且每遍历完一轮数据集才打印出一行更新，这也不足以让我们知道我们的训练到底怎样了。我们也许还应该：

• 训练过程中时不时地生成点图像样本，供我们检查模型性能
• 在训练过程中，把诸如损失值和生成的图片样本在内的一些东西记录到日志里。你可以使用诸如 Weights and Biases 或 tensorboard 之类的工具

我创建了一个快速的脚本程序(finetune_model.py)，使用了上述的训练代码并加入了少量日志记录功能。你可以在这里看看一次训练的日志：

%wandb johnowhitaker/dm_finetune/2upaa341 # You'll need a W&B account for this to work - skip if you don't want to log in
1

观察随着训练进展生成的样本图片如何变化也挺好玩 —— 即使从损失值看它好像并没有改进，但我们也能看到一个从原有图像分布（卧室图片）到新的数据集（wikiart 数据集）逐渐演变的过程。在这一节笔记本最后还有一些被注释掉的用于微调的代码，可以使用该脚本程序替代你运行上面的代码块。

# 练习： 看看你能不能修改第一单元的官方示例训练脚本程序
# 尝试使用预训练的模型，而不是从头开始训练
# 对比一下上面链接的最小化脚本 —— 对比一下哪些额外功能是最小化脚本没有的？
1
2
3

用这个模型生成点图片，我们可以看到这些脸看起来极其奇怪！

# @markdown Generate and plot some images:
x = torch.randn(8, 3, 256, 256).to(device)  # Batch of 8
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);
1
2
3
4
5
6
7
8
9

0it [00:00, ?it/s]
1

考虑因素4： 微调这个过程可能是难以预知的。如果我们训练很长时间，我们也许能看见一些生成得很完美的蝴蝶，但中间过程从模型自身讲也极其有趣，尤其是你对艺术风格感兴趣时！你可以试试短时间或长时间地观察一下训练过程，并试着该百年学习率，看看这会怎么影响模型的最终输出。

这个fine-tune跑了很久，结果果真出人意料，后面有空再尝试改变参数，后续保存到hugging face报了序列化错误，先做个笔记，再慢慢查原因。