1 实验介绍

1.1 实验目的

理解diffusion模型生成图像的主要步骤及原理
理解diffusion模型pytorch开发基本流程
掌握diffusion模型图像生成任务的pytorch实现详细流程
掌握diffusion模型典型应用场景及NPU迁移应用方法

1.2 实验环境

1.3 关于本实验

本实验使用pytorch框架，进行网络搭建、数据处理、模型训练和预测，完成MNIST手写体生成任务。

1.4 背景知识

pytorch基础知识，扩散模型背景知识，UNet网络encoder-decoder结构，diffusers模块的pipeline、噪声调度器等相关知识。

扩散模型（diffusion models）是深度生成模型中新的SOTA。扩散模型在图片生成任务中超越了原SOTA：GAN，并且在诸多应用领域都有出色的表现，如计算机视觉，NLP、波形处理、多模态建模、分子图建模、时间序列建模、对抗性净化等。扩散模型包含前向过程（扩散）和逆向过程（采样），前者不断给图片添加噪声，直到形成一张噪声图；后者从噪声图开始让模型学会拟合每一步的噪声分布并去噪，直到能恢复原始的图片。

1.5 实验设计

true

2 MNIST手写体生成实验

2.1 环境准备

torch模块主要用于本次实验扩散模型图像生成网络的构建:

torch.nn：主要包括网络可能涉及到的各类网络层
torch.utils.data：包括DataLoader子模块，用于MNIST数据集的载入预处理，也可以自定义数据集。

diffusers模块用于扩散模型优化网络调用，本实验主要涉及UNet2DModel子模块。在第3节NPU迁移实验中涉及不同任务pipeline及噪声调度器的使用。

matplotlib模块用于数据处理及模型训练过程可视化。

import torch 
import torchvision 
from torch import nn 
from torch.utils.data import DataLoader 
from diffusers import UNet2DModel 
from matplotlib import pyplot as plt

2.2 数据处理

定义数据预处理函数。函数功能包括：

加载数据集
数据集加噪
数据分步输出

代码：

加载数据集：

dataset = torchvision.datasets.MNIST(root="mnist/", train=True, download
 =True,transform=torchvision.transforms. ToTensor())
train_dataloader = DataLoader(dataset, batch_size=8,shuffle=True)
x, y = next(iter(train_dataloader))

数据集加噪：

def corrupt(x, amount):
    "根据amount为输入x加入噪声，这就是退化过程"
    noise = torch.rand_like(x)
    amount = amount.view(-1, 1, 1, 1) # 整理形状以保证广播机制不会出错
    return x * (1-amount) + noise*amount

amount = torch.linspace(0, 1, x.shape[0]) # 从0到1退化逐渐强烈
noised_x = corrupt(x, amount)

加噪图片预测结果： true

数据分步输出：

for i in range(n_steps):
    with torch.no_grad():                            # 在推理时不需要考虑张量的导数
        pred = net(x)          # 预测去噪后的图像
    pred_output_history.append(pred.detach().cpu())
    mix_factor = 1/(n_steps - i)   # 设置朝着预测方向移动多少
    x = x*(1-mix_factor) + pred*mix_factor
    step_history.append(x.detach().cpu())

训练20个epoch（左为每步输入图片，右为每步预测结果）： true

2.3 网络定义

参考UNet网络搭建扩散模型网络结构：Unet 网络结构是对称的，Unet 基于 Encoder-Decoder 结构，通过拼接的方式实现特征融合，结构简明且稳定。Encoder 负责特征提取，feature map 经过 Decoder 恢复原始分辨率。本实验将参考UNet论文，基于torch框架搭建简单的UNet网络：

class BasicNet(nn.Module):
    "一个简单的UNet网络"
    def __init__(self, in_channels=1, out_channles=1):
        super().__init__()
        self.down_layers = torch.nn.ModuleList([
            nn.Conv2d(in_channels, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
        ])
        self.up_layers = torch.nn.ModuleList([
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, out_channles, kernel_size=5, padding=2),
        ])
        self.act = nn.SiLU() # 激活函数
        self.downscale = nn.MaxPool2d(2)
        self.upscale = nn.Upsample(scale_factor=2)
        
    def forward(self, x):
        h = []
        for i, l in enumerate(self.down_layers):
            x = self.act(l(x))    # 通过运算层与激活函数
            if i < 2:     # 选择除了第3层（最后一层）以外的层
                h.append(x) # 排列供残差连接时使用的数据
                x = self.downscale(x) # 进行下采样以适配下一层的输入
                
        for i, l in enumerate(self.up_layers):
            if i > 0:    # 选择除了第1个上采样层以外的层
                x = self.upscale(x)  # Upscale上采样
                x += h.pop()    # 得到之前排列好的供残差连接使用的数据
            x = self.act(l(x)) # 通过运算层与激活函数            
        return x

或使用调用优化网络UNet2DModel的方式搭建扩散模型：

diffusers_model = UNet2DModel(
    sample_size = 28,
    in_channels = 1,
    out_channels = 1,
    layers_per_block=2,
    block_out_channels = (32, 64, 64),
    down_block_types=(
        "DownBlock2D",
        "AttnDownBlock2D",
        "AttnDownBlock2D",
    ),
    up_block_types=(
        "AttnUpBlock2D",
        "AttnUpBlock2D",
        "UpBlock2D",
    ),
)

2.4 模型训练

步骤1 加载数据集

见2.2节数据处理-加载数据集

步骤2 构建网络

构建网络，选择损失函数、优化器、模型

代码：

net = BasicNet()
x = torch.rand(8, 1, 28, 28)
net(x).shape
batch_size = 128
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
# 设置运行epoch
n_epoch = 70
# 创建网路
net = BasicNet()
net.to(device)
# 指定损失函数
loss_fn = nn.MSELoss()
# 指定优化器
opt = torch.optim.Adam(net.parameters(), lr=1e-3)
# 记录损失
losses = []

步骤3 训练模型

代码：

for epoch in range(n_epoch):
    for x, y in train_dataloader:
        # 得到数据并准备退化
        x = x.to(device)                          # 数据加载到GPU
        noise_amount = torch.rand(x.shape[0]).to(device)       # 随机选取噪声量
        noisy_x = corrupt(x, noise_amount)
        
        # 预测得到结果
        pred = net(noisy_x)
        
        # 计算损失
        loss = loss_fn(pred, x)  # 最重要的点：模型学到的是将一个噪声图片还原为mnist数字图片的能力
        
        # 反向传播并更新参数
        opt.zero_grad()
        loss.backward()
        opt.step()
        
        # 储存损失
        losses.append(loss.item())
        
    avg_loss = sum(losses[-len(train_dataloader):]) / len(train_dataloader)
    print(f"Finished epoch {epoch}. Average loss for this epoch:{avg_loss:05f}")

训练20个epoch结束时的打屏信息： true

2.5 推理预测

采样策略，把采样过程拆解为5步，每次只前进一步

n_steps = 40
x = torch.rand(64, 1, 28, 28).to(device) # 从完全随机的值开始
步骤2：分步预测结果
for i in range(n_steps):
    with torch.no_grad():               # 在推理时不需要考虑张量的导数
        pred = net(x)                           # 预测去噪后的图像
    pred_output_history.append(pred.detach().cpu())
    mix_factor = 1/(n_steps - i)            # 设置朝着预测方向移动多少
    x = x*(1-mix_factor) + pred*mix_factor   #每一步对输入叠加预测结果
    step_history.append(x.detach().cpu())    #更新每一步输入的加噪图片

训练70个epoch 每幅64个手写体图片的预测结果： true

3 NPU迁移实验

3.1 准备环境

当前模型支持的 PyTorch 版本和已知三方库依赖如下表所示。

表 1 版本支持表

Torch_Version 三方库依赖版本
PyTorch 1.11 diffusers==0.18.1 accelerate==0.20.3
环境准备指导。

请参考《Pytorch框架训练环境准备》搭建torch环境。

安装依赖。

在模型根目录下执行命令，安装模型对应PyTorch版本需要的依赖。

pip install -e .                    # 安装diffusers
cd examples/text_to_image/           # 根据下游任务安装对应依赖
pip install -r requirements.txt

修改三方库代码：

# 1. 在${python路径}/python3.7/site-packages/accelerate/accelerator.py中找到：
self.scaler = torch.cuda.amp.GradScaler(**kwargs)
修改为：
import torch_npu
self.scaler = torch.npu.amp.GradScaler(**kwargs)

# 2. 修改${python路径}/python3.7/site-packages/accelerate/utils/dataclasses.py，给类GradScalerKwargs添加属性：
dynamic: bool = True

3.2 推理实例

原例实现：https://gitee.com/ascend/ModelZoo-PyTorch/tree/master/PyTorch/built-in/diffusion/diffusers0.18.1

文生图

预训练模型准备：联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

CompVis/ldm-text2im-large-256

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

generator = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",torch_dtype=torch.float16)

修改prompt等操作需要对代码进行修改设置图像生成的文本内容：

image = generator("The pastoral scenery in classical style
").images[0]

运行在线推理：

python test_infer/text-to-image.py

生成效果：古典油画风格的田园风景 true

文本指导图生图

预训练模型准备：联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

nitrosocke/Ghibli-Diffusion

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

pipe = StableDiffusionImg2ImgPipeline.from_pretrained("nitrosocke/Ghibli-Diffusion").to(device)

设置图像生成的内容：

prompt = "Monet style, a fantasy landscape with mountain stream"

修改test_infer/text-guide-img-to-img.py中url为本地图片地址：

url = "noise.JPG"

true

设置prompt：

prompt = "Monet style, a fantasy landscape with mountain stream"

运行推理代码：

python test_infer/text-guide-img-to-img.py

图像生成效果：莫奈风格的溪流风景画 true

文本指导图像修复

预训练模型准备：联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

runwayml/stable-diffusion-inpainting

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

pipeline = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")

设置原图的内容：

true

设置原图mask处理后的内容：

true

修改prompt：

prompt = "A white cute dog leaning against a railing"

运行推理代码：

python test_infer/text-guide-image-inpainting.py

图像修复效果：掩膜部分替换为小狗倚靠栅栏，生成的内容较好地融入了背景部分

true

文本指导图像深度生成

预训练模型准备：联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

stabilityai/stable-diffusion-2-depth

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

pipe = StableDiffusionDepth2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-depth").to("npu")

设置原图输入：

true

设置prompt：

prompt = "a snowflake"
n_prompt = "blue, clear, serene, crystal like"

运行推理代码：

python test_infer/text-guide-depth-to-image.py

图像深度生成效果：保留原始图片纹理色彩的基础上生成了雪花图片

true

无条件图像生成

预训练模型准备：联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

anton-l/ddpm-butterflies-128

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

generator = DiffusionPipeline.from_pretrained("anton-l/ddpm-butterflies-128")

运行在线推理：

python test_infer/unconditional-image-generation.py

图像生成效果：按照预训练的模型设置生成了蝴蝶图片

true

4 实验总结与思考

4.1 实验总结

本实验介绍了pytorch在图像生成上的应用，使用MNIST手写体数据集搭建UNet神经网络，基于扩散模型原理进行手写体图像生成，让学员熟悉diffusion模型的pytorch实现流程，掌握diffusion模型的应用场景及NPU迁移应用方法。

4.2 思考题

diffusion模型图像生成主要包括哪些步骤？

前向扩散（加噪）-> 训练不同图片添加噪声的分布 -> 采样还原图片（去噪）
DDIM与DDPM各自的优势？

DDPM适用于生成模型的稳定和准确训练，DDIM专注于从训练好的模型中高效生成样本。
将1个batch的图片合成在1张图中使用哪个函数？

torchvision.utils.make_grid()
向原始图片加入逐渐增大的噪声可使用什么函数？

torch.linspace()、torch.sigmoid()、DDPMScheduler.add_noise()
Diffusers核心API主要分为哪几个部分：

管线（DDPMPipeline）、模型、调度器（DDPMScheduler）

5 拓展实验

图像加噪训练及生成过程中，可尝试对扩散模型进行适当改进及训练策略优化，（DDPM、DDIM或其他方式进行预测？保留一些历史预测值指导当前步更新？等）提升MNIST手写体生成任务的呈现效果。
根据不同数据集特点微调模型进行迁移以及添加文字引导不同的图像生成实验。
合理使用噪声调度器，调节每步噪声添加的幅度，探索最佳噪声量实验。

实验	介绍	难度	软件环境	开发环境
MNIST手写体生成实验	基于pytorch框架，使用MNIST数据集搭建diffusers神经网络UNet及UNet2DModel, 进行模型训练及预测	简单	python3.7、pytorch1.11	PC机

1 实验介绍

1.1 实验目的

1.2 实验环境

1.3 关于本实验

1.4 背景知识

1.5 实验设计

2 MNIST手写体生成实验

2.1 环境准备

2.2 数据处理

2.3 网络定义

2.4 模型训练

2.5 推理预测

3 NPU迁移实验

3.1 准备环境

3.2 推理实例

文生图

文本指导图生图

文本指导图像修复

文本指导图像深度生成

无条件图像生成

4 实验总结与思考

4.1 实验总结

4.2 思考题

5 拓展实验

关于昇腾

新闻与活动

交流与资讯

支持与服务

开源社区

Torch_Version	三方库依赖版本
PyTorch 1.11	diffusers==0.18.1 accelerate==0.20.3