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函數在一個大圖中展示了所有步驟的圖像。
)
