devbox
知新_RL
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

如何训练一个简单的stable diffusion模型(附详细注释）_stable diffusion如何训练模型

作者：知新_RL | 2024-05-29 14:39:23

赞

踩

stable diffusion如何训练模型

注：代码来自https://github.com/darcula1993/diffusion-models-class-CN/blob/main/unit1/01_introduction_to_diffusers_CN.ipynb 

本文是本人学习后的的尝试以及注解

一、准备工作

  1. """
  3. 这行命令使用pip工具来安装或升级多个Python包。具体来说，它执行以下操作：
  5. -qq：这是pip的安静模式选项，它会减少输出信息，只显示关键信息，使安装过程更为简洁。
  7. -U：这是pip的升级选项，它指示pip升级已经安装的包到最新版本（如果存在新版本）。
  9. 接下来，列出了要安装或升级的包：
  11. diffusers：一个Python包
  13. datasets：Hugging Face Transformers库的一部分，用于提供和管理各种自然语言处理（NLP）数据集的工具。
  15. transformers：Hugging Face Transformers库，提供了预训练的深度学习模型，用于自然语言处理和文本生成任务。
  17. accelerate：Hugging Face库的一部分，用于加速深度学习模型的训练和推理。
  19. ftfy：一个用于处理Unicode文本的Python库，用于修复和清理不规范的Unicode文本。
  21. pyarrow：正如前面所提到的，pyarrow是一个用于高效处理大规模数据集的Python库，支持列式存储和跨语言互操作性。
  22. """
  23.  
  24. %pip install -qq -U diffusers datasets transformers accelerate ftfy pyarrow
  1. # 登录hugging face
  2. from huggingface_hub import notebook_login
  3.  
  4. notebook_login()

显示下图则登陆成功：

  1. # 安装 Git LFS 来上传模型检查点：
  2.  
  3. %%capture
  4. !sudo apt -qq install git-lfs
  5. !git config --global credential.helper store
  1. # 导入将要使用的库，并定义一些方便函数，稍后将会使用这些函数
  2. import numpy as np
  3. import torch
  4. import torch.nn.functional as F
  5. from matplotlib import pyplot as plt
  6. from PIL import Image
  7.  
  8.  
  9. def show_images(x):
  10.   """Given a batch of images x, make a grid and convert to PIL"""
  11.   '''
  12.   输入参数：x，一个批量的图像数据（通常是PyTorch张量）。
  13.   功能：将输入的图像数据从范围(-1, 1)映射到(0, 1)，然后将这些图像排列成一个网格，并将网格转换为PIL图像。
  14.   返回值：返回一个PIL图像，其中包含了排列好的输入图像网格。
  15.   '''
  16.   x = x * 0.5 + 0.5  # Map from (-1, 1) back to (0, 1)
  17.   grid = torchvision.utils.make_grid(x)
  18.   grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
  19.   grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
  20.   return grid_im
  21.  
  22.  
  23. def make_grid(images, size=64):
  24.   """Given a list of PIL images, stack them together into a line for easy viewing"""
  25.   '''
  26.   输入参数：images，一个包含多个PIL图像的列表，以及一个可选的size参数，用于指定图像的大小。
  27.   功能：将多个PIL图像按照指定的大小堆叠在一行上，以便于查看。
  28.   返回值：返回一个新的PIL图像，其中包含了堆叠在一行上的输入图像。
  29.   '''
  30.   output_im = Image.new("RGB", (size * len(images), size))
  31.   for i, im in enumerate(images):
  32.       output_im.paste(im.resize((size, size)), (i * size, 0))
  33.   return output_im
  34.  
  35.  
  36. # Mac users may need device = 'mps' (untested)
  37. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

二、下载训练数据集

  1. # 下载一个来自 Hugging Face Hub 的图像集。具体来说，是个 1000 张蝴蝶图像收藏集
  2. import torchvision  # 导入PyTorch的torchvision库
  3. from datasets import load_dataset  # 导入数据集
  4. from torchvision import transforms  # 导入PyTorch的transforms模块
  5.  
  6. # 使用load_dataset函数加载名为"huggan/smithsonian_butterflies_subset"的数据集中的训练集数据
  7. dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")
  8.  
  9. # 或者从本地文件夹加载图像数据
  10. # dataset = load_dataset("imagefolder", data_dir="path/to/folder")
  11.  
  12. # 指定图像大小为32x32像素
  13. image_size = 32
  14. # 如果GPU内存不足，可以降低批次大小
  15. batch_size = 64
  16.  
  17. # 定义数据增强操作
  18. preprocess = transforms.Compose(
  19.     [
  20.         transforms.Resize((image_size, image_size)),  # 调整图像大小
  21.         transforms.RandomHorizontalFlip(),  # 随机水平翻转（数据增强）
  22.         transforms.ToTensor(),  # 将图像转换为张量（数值范围从0到1）
  23.         transforms.Normalize([0.5], [0.5]),  # 将图像像素值归一化到(-1, 1)的范围
  24.     ]
  25. )
  26.  
  27. # 定义一个用于对数据进行转换的函数
  28. def transform(examples):
  29.     images = [preprocess(image.convert("RGB")) for image in examples["image"]]
  30.     return {"images": images}
  31.  
  32. # 将数据集的转换函数设置为上述定义的transform函数
  33. dataset.set_transform(transform)
  34.  
  35. # 创建一个数据加载器，用于以批次方式提供转换后的图像数据
  36. train_dataloader = torch.utils.data.DataLoader(
  37.     dataset, batch_size=batch_size, shuffle=True  # 使用指定的批次大小和随机打乱数据
  38. )
  1. # 我们可以从中取出一批图像数据来看一看他们是什么样子:
  2.  
  3. xb = next(iter(train_dataloader))["images"].to(device)[:8]
  4. print("X shape:", xb.shape)
  5. show_images(xb).resize((8 * 64, 64), resample=Image.NEAREST)

结果如下图所示：

三、定义管理器

我们的训练计划是，取出这些输入图片然后对它们增添噪声，在这之后把带噪的图片送入模型。在推理阶段，我们将用模型的预测值来不断迭代去除这些噪点。在diffusers中，这两个步骤都是由 管理器（调度器） 来处理的。噪声管理器决定在不同的迭代周期时分别加入多少噪声。

  1. from diffusers import DDPMScheduler
  2.  
  3. noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
  1. # 绘图查看输入 (x) 与噪声是如何在不同迭代周期中量化和叠加的
  2. plt.plot(noise_scheduler.alphas_cumprod.cpu() ** 0.5, label=r"${\sqrt{\bar{\alpha}_t}}$")
  3. plt.plot((1 - noise_scheduler.alphas_cumprod.cpu()) ** 0.5, label=r"$\sqrt{(1 - \bar{\alpha}_t)}$")
  4. plt.legend(fontsize="x-large");

  1. # 使用noise_scheduler.add_noise功能来添加不同程度的噪声
  2. timesteps = torch.linspace(0, 999, 8).long().to(device)
  3. noise = torch.randn_like(xb)
  4. noisy_xb = noise_scheduler.add_noise(xb, noise, timesteps)
  5. print("Noisy X shape", noisy_xb.shape)
  6. show_images(noisy_xb).resize((8 * 64, 64), resample=Image.NEAREST)

四、定义、训练模型

  1. # 定义模型
  2. from diffusers import UNet2DModel
  3.  
  4. # Create a model
  5. model = UNet2DModel(
  6.   sample_size=image_size,  # the target image resolution
  7.   in_channels=3,  # the number of input channels, 3 for RGB images
  8.   out_channels=3,  # the number of output channels
  9.   layers_per_block=2,  # how many ResNet layers to use per UNet block
  10.   block_out_channels=(64, 128, 128, 256),  # More channels -> more parameters
  11.   down_block_types=(
  12.       "DownBlock2D",  # a regular ResNet downsampling block
  13.       "DownBlock2D",
  14.       "AttnDownBlock2D",  # a ResNet downsampling block with spatial self-attention
  15.       "AttnDownBlock2D",
  16.   ),
  17.   up_block_types=(
  18.       "AttnUpBlock2D",
  19.       "AttnUpBlock2D",  # a ResNet upsampling block with spatial self-attention
  20.       "UpBlock2D",
  21.       "UpBlock2D",  # a regular ResNet upsampling block
  22.   ),
  23. )
  24. model.to(device);

开始训练模型 

  1. # 创建了一个对象，用于调度噪声的添加。参数指定了训练的总步数和噪声的变化规律
  2. noise_scheduler = DDPMScheduler(
  3.     num_train_timesteps=1000, beta_schedule="squaredcos_cap_v2"
  4. )
  5.  
  6. # 创建了一个AdamW优化器，用于更新模型的参数。返回模型的可训练参数，参数指定了学习率。
  7. optimizer = torch.optim.AdamW(model.parameters(), lr=4e-4)
  8.  
  9. # 创建一个空列表，用于存储每个步骤的损失值。
  10. losses = []
  11.  
  12. for epoch in range(30):
  13.   for step, batch in enumerate(train_dataloader):
  14.     # 从批次中获取干净的图像数据，并将其移动到指定的设备（例如GPU）上进行计算。
  15.     clean_images = batch["images"].to(device)
  16.     # 生成与干净图像相同形状的噪声张量，该噪声将被添加到图像中。
  17.     noise = torch.randn(clean_images.shape).to(clean_images.device)
  18.     # 获取批次的大小。
  19.     bs = clean_images.shape[0]
  20.  
  21.     # 为每个图像随机生成一个时间步长，该时间步长将用于确定噪声的变化程度。
  22.     timesteps = torch.randint(
  23.       0, noise_scheduler.num_train_timesteps, (bs,), device=clean_images.device
  24.     ).long()
  25.  
  26.     # 根据噪声调度器中的每个噪声幅度和时间步长，将噪声添加到干净图像中，生成带有噪声的图像。
  27.     noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
  28.  
  29.     # 使用模型对带有噪声的图像进行预测，得到去噪后的图像。
  30.     noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
  31.  
  32.     # 计算预测图像与真实噪声之间的均方误差损失，计算损失相对于模型参数的梯度，并将当前步骤的损失添加到列表中。
  33.     loss = F.mse_loss(noise_pred, noise)
  34.     loss.backward(loss)
  35.     losses.append(loss.item())
  36.  
  37.     # 使用优化器更新模型的参数，并将梯度置零。
  38.     optimizer.step()
  39.     optimizer.zero_grad()
  40.  
  41.   # 每隔5个epoch，计算最近一个epoch的平均损失，并打印出来。
  42.   if (epoch + 1) % 5 == 0:
  43.       loss_last_epoch = sum(losses[-len(train_dataloader) :]) / len(train_dataloader)
  44.       print(f"Epoch:{epoch+1}, loss: {loss_last_epoch}")

  1. # 绘制 loss 曲线，我们能看到模型在一开始快速的收敛，接下来以一个较慢的速度持续优化（我们用右边 log 坐标轴的视图可以看的更清楚）：
  2.  
  3. fig, axs = plt.subplots(1, 2, figsize=(12, 4))
  4. axs[0].plot(losses)
  5. axs[1].plot(np.log(losses))
  6. plt.show()

五、生成图像 

下面开始生成图像 

  1. # 方法 1：建立一个管道：
  2. from diffusers import DDPMPipeline
  3.  
  4. image_pipe = DDPMPipeline(unet=model, scheduler=noise_scheduler)
  5. pipeline_output = image_pipe()
  6. pipeline_output.images[0]

  1. # 我们可以在本地文件夹这样保存一个管道：
  2. image_pipe.save_pretrained("my_pipeline")
  3.  
  4. # 检查文件夹的内容：
  5. !ls my_pipeline/

 这里scheduler与unet子文件夹中包含了生成图像所需的全部组件。比如，在unet文件中能看到模型参数 (diffusion_pytorch_model.bin) 与描述模型结构的配置文件。

  1. # 方法 2：写一个取样循环
  2.  
  3. # 从随机噪声开始，遍历管理器的迭代周期来看从最嘈杂直到最微小的噪声变化，基于模型的预测一步步减少一些噪声：
  4.  
  5. # Random starting point (8 random images):
  6. sample = torch.randn(8, 3, 32, 32).to(device)
  7.  
  8. for i, t in enumerate(noise_scheduler.timesteps):
  9.  
  10.     # Get model pred
  11.     with torch.no_grad():
  12.         residual = model(sample, t).sample
  13.  
  14.     # Update sample with step
  15.     sample = noise_scheduler.step(residual, t, sample).prev_sample
  16.     
  17. # noise_scheduler.step () 函数相应做了 sample（取样）时的数学运算。
  18. show_images(sample)

 六、将模型上传到Hugging Face

在上面的例子中我们把管道保存在了本地。 把模型 push 到 hub 上，我们会需要建立模型和相应文件的仓库名。 我们根据你的选择（模型 ID）来决定仓库的名字（大胆的去替换掉model_name吧；需要包含你的用户名，get_full_repo_name ()会帮你做到）：

  1. from huggingface_hub import get_full_repo_name
  2.  
  3. model_name = "sd-class-butterflies-32"
  4. hub_model_id = get_full_repo_name(model_name)
  5. hub_model_id
  1. # 然后，在 
    声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/知新_RL/article/detail/642663
    推荐阅读
    相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。