DeepSeek基于注意力模型的可控图像生成

DeepSeek大模型高性能核心技术与多模态融合开发 - 商品搜索 - 京东

图像的加噪与模型训练

在扩散模型的训练过程中，首先需要对输入的信号进行加噪处理，经典的加噪过程是在图像进行向量化处理后在其中添加正态分布，而正态分布的值也是与时间步相关的。这样逐步向图像中添加噪声，直到图像变得完全噪声化。

import torch   T = 1000  # Diffusion过程的总步数  # 前向diffusion计算参数
# (T,) 生成一个线性间隔的tensor，用于计算每一步的噪声水平  
betas = torch.linspace(0.0001, 0.02, T)    
alphas = 1 - betas  # (T,) 计算每一步的保留率  
# alpha_t累乘 (T,)，计算每一步累积的保留率 
alphas_cumprod = torch.cumprod(alphas, dim=-1)   
# alpha_t-1累乘(T,)，为计算方差做准备
alphas_cumprod_prev = torch.cat((torch.tensor([1.0]), alphas_cumprod[:-1]), dim=-1)    
# denoise用的方差(T,)，计算每一步的去噪方差 
variance = (1 - alphas) * (1 - alphas_cumprod_prev) / (1 - alphas_cumprod)   # 执行前向加噪  
def forward_add_noise(x, t):  # batch_x: (batch,channel,height,width), batch_t: (batch_size,)  noise = torch.randn_like(x)  # 为每幅图片生成第t步的高斯噪声   (batch,channel,height,width)  # 根据当前步数t，获取对应的累积保留率，并调整其形状以匹配输入x的形状    batch_alphas_cumprod = alphas_cumprod[t].view(x.size(0), 1, 1, 1)    # 基于公式直接生成第t步加噪后的图片    x = torch.sqrt(batch_alphas_cumprod) * x + torch.sqrt(1 - batch_alphas_cumprod) * noise    return x, noise  # 返回加噪后的图片和生成的噪声

上面这段代码首先定义了扩散模型的前向过程中需要的参数，包括每一步的噪声水平betas、保留率alphas、累积保留率alphas_cumprod以及用于去噪的方差variance。然后定义了一个函数forward_add_noise，该函数接受一个图像x和步数t作为输入。根据扩散模型的前向过程，向图像中添加噪声，并返回加噪后的图像和生成的噪声。

读者可以采用以下代码尝试完成为图像添加噪声的演示：

import matplotlib.pyplot as plt 
from dataset import MNISTdataset=MNIST()
# 两幅图片拼batch, (2,1,48,48)    
x=torch.stack((dataset[0][0],dataset[1][0]),dim=0) # 原图
plt.figure(figsize=(10,10))
plt.subplot(1,2,1)
plt.imshow(x[0].permute(1,2,0))
plt.subplot(1,2,2)
plt.imshow(x[1].permute(1,2,0))
plt.show()# 随机时间步
t=torch.randint(0,T,size=(x.size(0),))
print('t:',t)# 加噪
x=x*2-1 # [0,1]像素值调整到[-1,1]之间，以便与高斯噪声值范围匹配
x,noise=forward_add_noise(x,t)
print('x:',x.size())
print('noise:',noise.size())# 加噪图
plt.figure(figsize=(10,10))
plt.subplot(1,2,1)
plt.imshow(((x[0]+1)/2).permute(1,2,0))   
plt.subplot(1,2,2)
plt.imshow(((x[0]+1)/2).permute(1,2,0))
plt.show()

运行结果如图9-13所示。

在此基础上，我们可以完成对Dit模型的训练，代码如下：

from torch.utils.data import DataLoader  # 导入PyTorch的数据加载工具  
from dataset import MNIST  # 从dataset模块导入MNIST数据集类  
from diffusion import forward_add_noise  # 从diffusion模块导入forward_add_noise函数，用于向图像添加噪声  
import torch  # 导入PyTorch库  
from torch import nn  # 从PyTorch导入nn模块，包含构建神经网络所需的工具  
import os  # 导入os模块，用于处理文件和目录路径  
from dit import DiT  # 从dit模块导入DiT模型  
# 判断是否有可用的CUDA设备，如果有则使用GPU，否则使用CPU  
DEVICE='cuda' if torch.cuda.is_available() else 'cpu'    dataset=MNIST()  # 实例化MNIST数据集对象  T = 1000  # 设置扩散过程中的总时间步数  
model=DiT(img_size=28,patch_size=4,channel=1,emb_size=64,label_num=10,dit_num=3,head=4).to(DEVICE)  # 实例化DiT模型并移至指定设备  
#model.load_state_dict(torch.load('./saver/model.pth'))  # 可选：加载预训练模型参数  # 使用Adam优化器，学习率设置为0.001
optimzer=torch.optim.Adam(model.parameters(),lr=1e-3)   
loss_fn=nn.L1Loss()  # 使用L1损失函数（即绝对值误差均值）  '''训练模型'''  
EPOCH=300  # 设置训练的总轮次  
BATCH_SIZE=300  # 设置每个批次的大小  if __name__ == '__main__':  from tqdm import tqdm  # 导入tqdm库，用于在训练过程中显示进度条  dataloader=DataLoader(dataset,batch_size=BATCH_SIZE,shuffle=True,num_workers=10,persistent_workers=True)  # 创建数据加载器  iter_count=0  for epoch in range(EPOCH):  # 遍历每个训练轮次  pbar = tqdm(dataloader, total=len(dataloader))  # 初始化进度条  for imgs,labels in pbar:  # 遍历每个批次的数据  x=imgs*2-1  # 将图像的像素范围从[0,1]转换到[-1,1]，与噪声高斯分布的范围对应  t=torch.randint(0,T,(imgs.size(0),))  # 为每幅图片生成一个随机的t时刻  y=labels  # 向图像添加噪声，返回加噪后的图像和添加的噪声x,noise=forward_add_noise(x,t)    # 模型预测添加的噪声pred_noise=model(x.to(DEVICE),t.to(DEVICE),y.to(DEVICE))    # 计算预测噪声和实际噪声之间的L1损失loss=loss_fn(pred_noise,noise.to(DEVICE))    optimzer.zero_grad()  # 清除之前的梯度  loss.backward()  # 反向传播，计算梯度  optimzer.step()  # 更新模型参数  # 更新进度条描述pbar.set_description(f"epoch:{epoch + 1}, train_loss:{loss.item():.5f}")    if epoch % 20 == 0:  # 每20轮保存一次模型  torch.save(model.state_dict(),'./saver/model.pth')  print("base diffusion saved")

 读者自行查看代码运行结果。

基于注意力模型的可控图像生成

DiT模型的可控图像生成是在我们训练的基础上，逐渐对正态分布的噪声图像进行按步骤的脱噪过程。这一过程不仅要求模型具备精准的噪声预测能力，还需确保脱噪步骤的细腻与连贯，从而最终实现从纯粹噪声到目标图像的华丽蜕变。

完整的可控图像生成代码如下：

import torch   from dit import DiT  
import matplotlib.pyplot as plt   
# 导入diffusion模块中的所有内容，这通常包含一些与扩散模型相关的预定义变量和函数
from diffusion import *    # 设置设备为GPU或CPU  
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'   
DEVICE = "cpu"  # 强制使用CPU  T = 1000  # 扩散步骤的总数  def backward_denoise(model,x,y):  steps=[x.clone(),]  # 初始化步骤列表，包含初始噪声图像  global alphas,alphas_cumprod,variance  # 这些是从diffusion模块导入的全局变量  x=x.to(DEVICE)  # 将输入x移动到指定的设备  alphas=alphas.to(DEVICE)  alphas_cumprod=alphas_cumprod.to(DEVICE)  variance=variance.to(DEVICE)  y=y.to(DEVICE)  # 将标签y移动到指定的设备  model.eval()  # 设置模型为评估模式  with torch.no_grad():  # 在不计算梯度的情况下运行，节省内存和计算资源  for time in range(T-1,-1,-1):  # 从T-1到0逆序迭代  t=torch.full((x.size(0),),time).to(DEVICE)  # 创建一个包含当前时间步的tensor  # 预测x_t时刻的噪声  noise=model(x,t,y)      # 生成t-1时刻的图像  shape=(x.size(0),1,1,1)  mean=1/torch.sqrt(alphas[t].view(*shape))*  \  (  x-   (1-alphas[t].view(*shape))/torch.sqrt(1-alphas_cumprod[t].view(*shape))*noise  )  if time!=0:  x=mean+ \  torch.randn_like(x)* \  torch.sqrt(variance[t].view(*shape))  else:  x=mean  x=torch.clamp(x, -1.0, 1.0).detach()  # 确保x的值在[-1,1]之间，并分离计算图  steps.append(x)  return steps  # 初始化DiT模型  
model=DiT(img_size=28,patch_size=4,channel=1,emb_size=64,label_num=10,dit_num=3,head=4).to(DEVICE)  
model.load_state_dict(torch.load('./saver/model.pth'))  # 加载模型权重  # 生成噪声图  
batch_size=10  
x=torch.randn(size=(batch_size,1,28,28))  # 生成随机噪声图像  
y=torch.arange(start=0,end=10,dtype=torch.long)   # 生成标签  # 逐步去噪得到原图  
steps=backward_denoise(model,x,y)  # 绘制数量  
num_imgs=20  # 绘制还原过程  
plt.figure(figsize=(15,15))  
for b in range(batch_size):  for i in range(0,num_imgs):  idx=int(T/num_imgs)*(i+1)  # 计算要绘制的步骤索引  # 像素值还原到[0,1]  final_img=(steps[idx][b].to('cpu')+1)/2  # tensor转回PIL图  final_img=final_img.permute(1,2,0)  # 调整通道顺序以匹配图像格式  plt.subplot(batch_size,num_imgs,b*num_imgs+i+1)plt.imshow(final_img)  
plt.show()  # 显示图像

上面的代码展示了使用DiT进行图像去噪的完整过程。首先，它导入了必要的库和模块，包括PyTorch、DiT模型、matplotlib绘图模块，以及从diffusion模块导入的一些预定义变量和函数，这些通常与扩散模型相关。然后，代码设置了计算设备为CPU（尽管提供了检测GPU可用性的选项），并定义了扩散步骤的总数。

backward_denoise函数是实现图像去噪的核心。它接受一个DiT模型、一批噪声图像以及对应的标签作为输入。在这个函数内部，它首先将输入移动到指定的计算设备，然后将模型设置为评估模式，并开始一个不计算梯度的循环，从最后一个扩散步骤开始逆向迭代至第一步。在每一步中，它使用模型预测当前步骤的噪声，然后根据扩散模型的公式计算上一步的图像。这个过程一直持续到生成原始图像。

接下来，代码初始化了DiT模型，并加载了预训练的权重。然后，它生成了一批随机噪声图像和对应的标签，并使用backward_denoise函数对这些噪声图像进行去噪，逐步还原出原始图像。运行结果如图9-14所示。

图9-14  基于DiT模型的可控图像生成

可见，我们使用生成代码绘制了去噪过程的图像，展示了从完全噪声的图像逐步还原为清晰图像的过程。通过调整通道顺序和像素值范围，它将tensor格式的图像转换为适合绘制的格式，并使用matplotlib库的subplot函数在一个大图中展示了所有步骤的图像。