前言
本期文章,我們一起來探究生成式 AI 這一火熱的新知識領域。
目前計劃有三個大方向:
-
代碼深度實踐方向。例如用代碼完整诠釋 Diffusion 模型的工作原理,或者 Transformer 的完整架構等;
-
模型部署和訓練優化方向。例如嘗試解讀 LMI、DeepSpeed、Accelerate、FlashAttention 等不同模型優化方向的最新進展;
-
模型量化實踐方向。例如 GPTQ、bitsandbtyes 等前沿模型量化原理和實踐等。
在之前文章中,除了通過論文介紹生成式 AI 和大語言模型(LLMs)的主要原理之外,在代碼實踐環節主要還是局限于是引入預訓練模型、在預訓練模型基礎上做微調、使用 API 等等。很多資深研究者通過多種渠道和我們溝通,覺得還不過瘾,希望内容可以更加深入。
因此,本期做爲代碼深度實踐方向的第一個系列:“擴散模型原理”代碼實踐系列,将嘗試用代碼完整從底層開始洞悉擴散模型(Diffusion Models)的工作原理。而不再僅僅止步于引入預訓練模型或使用 API 完成工作。
擴散模型系列内容概述
基于擴散模型(Diffusion Models)的大模型,例如:Stable Diffusion、Midjourney、DALL-E 等能夠僅通過提示詞(Prompt)就能夠生成圖像。我們希望通過編寫這個“擴散模型原理”代碼實踐系列,使用代碼來探究和诠釋這些應用背後算法的原理。
這個由四篇文章組成的亞馬遜雲科技“擴散模型原理” 代碼實踐系列中,我們将:
-
探索基于擴散的生成人工智能的前沿世界,并從頭開始創建自己的擴散模型
-
深入了解擴散過程和驅動擴散過程的模型,而不僅僅是預先構建的模型和 API
-
在整個系列的最後,我們将有一個模型,可以作爲我們繼續探索應用擴散模型的起點
我将會用四集的篇幅,逐行代碼來構建擴散模型(Diffusion Model)。這四部分分别是:
-
噪聲采樣(Sampling)
-
訓練擴散模型(Training)
-
添加上下文(Embedding & Adding Context)
-
噪聲快速采樣(Fast Sampling)
這四部分的完整代碼可參考:https://github.com/hanyun2019/difussion-model-code-implementation
本文是第一部分:噪聲采樣(Sampling)。
擴散模型的目标
中國有句古語:起心動念。因此,既然我們要開始從底層揭開擴散模型(Diffusion Model)的面紗,首先是否應該要想清楚一個問題:使用擴散模型的目标是什麽?
本章将讨論擴散模型的目标,以及如何利用各種遊戲角色圖像(例如:精靈圖像)訓練數據來增強模型的能力,然後讓擴散模型自己去生成更多的遊戲角色圖像(例如:生成某種風格的精靈圖像等)。
假設下面是你已經有的精靈圖像數據集(來自 ElvGames 的 FrootsnVeggies 和 kyrise 精靈圖像集),你想要更多的在這些數據集中沒有的大量精靈圖像,你該如何實現這個現在看起來不可能完成的任務?
-
《FrootsnVeggies》
https://zrghr.itch.io/froots-and-veggies-culinary-pixels
-
《kyrise》
https://kyrise.itch.io/
Source: Sprites by ElvGames
面對這個看上去不可能完成的任務,擴散模型(Diffusion Model)就能幫上忙了。你有很多訓練數據,比如你在這裏看到的遊戲中精靈角色的圖像,這是你的訓練數據集。而你想要更多訓練數據集中沒有的精靈圖像。你可以使用神經網絡,按照擴散模型過程爲你生成更多這樣的精靈。擴散模型能夠生成這樣的精靈圖像。這就是我們這個系列要讨論的有趣話題。
以這個精靈圖像數據集爲例,擴散模型能夠學習到精靈角色的通用特征,例如某種精靈的身體輪廓、頭發顔色甚至腰帶配飾細節等。
神經網絡學習生成精靈圖像的概念是什麽呢?它可能是一些精緻的細節,比如精靈的頭發顔色、腰帶配飾等;也可能是一些大緻的輪廓,比如頭部輪廓、身體輪廓、或者介于兩者之間的其它輪廓等。而做到這一點的一種方法,即通過獲取數據并能夠專注更精細的細節或輪廓的方法,實際上是添加不同級别的噪聲(noise)。因此,這隻是在圖像中添加噪聲,它被稱爲 “噪聲過程”(noising process)。
這個思路其實是受到了物理學的啓發,場景很類似一滴墨水滴到一杯清水裏的全過程。最初我們确切地知道墨水滴落在那裏;但是随着時間的推移,我們會看到墨水擴散到清水中直到它完全消失(或者說完全和清水融爲一體)。
如下圖所示,我們從最左邊的圖像“Bob the Sprite”開始,當添加噪音時,它會消失,直到我們辨别不出它到底是哪個精靈。
Source: How Diffusion Models Work,https://learn.deeplearning.ai/diffusion-models/lesson/2/intuition ,by DeepLearning.AI
以這個 Bob 精靈圖像爲例,以下詳細描述通過添加不同階段噪聲,到精靈訓練數據集的全過程。
在最左邊圖像“Bob the Sprite!”的時候,我們想讓神經網絡知道:“這就是 Bob ,它是一個精靈”。
到了“Probably Bob”的時候,我們想讓神經網絡知道:“你知道,這裏有一些噪聲”,不過通過一些細節它看起來像“Bob the Sprite!”。
到了“Well, Bob or Fred”這個圖像時,變得隻能看到精靈的模糊輪廓了。那麽在這裏我們感覺到這可能是精靈,但可能是精靈 Bob 、精靈 Fred ,或者是精靈 Nance ,這時我們可能想讓神經網絡爲這些精靈圖像推薦更通用的細節,比如:在此基礎上爲 Bob 建議一些細節,或者你會爲 Fred 建議一些細節等。
到了最後“No Idea”這個圖像時,雖然已經無法辨認圖像的特征,我們仍然希望它看起來更像精靈。這時,我們仍然想讓神經網絡知道:“我希望你通過這張完全嘈雜的圖像,通過提煉出精靈可能樣子的輪廓,來把它變成更像精靈的圖像”。
這就是整個“噪聲過程”(noising process),即随着時間的推移逐漸增加噪聲的過程,如同把一滴墨水完全擴散到一杯清水之中。我們需要訓練的那個神經網絡,就是希望它能夠把不同的嘈雜圖像變成美麗精靈。這就是我們的目标,即擴散模型的目标。
要讓神經網絡做到這一點,就是要讓它學會去除添加的噪聲。從“No Idea”這個圖像開始(這時隻是純粹的噪聲),到開始看起來像裏面可能有精靈,再到長得像精靈 Bob ,到最後就是精靈 Bob。
這裏要強調的是:“No Idea”這個圖像的噪聲非常重要,因爲它是正态分布(normal distribution)的。換句話說,也就是這個圖像的像素每一個都是從正态分布(又稱 “高斯分布”)中采樣的。
因此,當你希望神經網絡生成一個新的精靈時,比如精靈 Fred ,你可以從該正态分布中采樣噪聲,然後你可以使用神經網絡逐漸去除噪聲來獲得一個全新的精靈!除了你訓練過的所有精靈之外,你還可以獲得更多的精靈。
Source:How Diffusion Models Work, https://learn.deeplearning.ai/diffusion-models/lesson/2/intuition ,by DeepLearning.AI
恭喜你,你已經找到了生成大量的全新美麗精靈的理論方法!接下來就是代碼實踐了。
在下一章裏,我們将用代碼展示爲了實現正态分布噪聲采樣,而主動在叠代階段添加噪聲的方法;和沒有添加噪聲方法的模型輸出結果對比測試。這将是一次很有趣和難忘的擴散模型工作原理奇妙體驗。
噪聲采樣的代碼實踐
首先我們将讨論采樣。我們将詳細介紹采樣的細節以及它在多個不同的叠代中是如何工作的。
1. 創建 Amazon SageMaker Notebook 實例
篇幅所限,本文不再贅述如何創建 Amazon SageMaker Notebook 實例。
如需詳細了解,可參考以下官方文檔:
https://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/gs-setup-working-env.html
2. 代碼說明
本實驗的完整示例代碼可參考:https://github.com/hanyun2019/difussion-model-code-implementation/blob/dm-project-haowen-mac/L1_Sampling.ipynb
示例代碼的 notebook 在 Amazon SageMaker Notebook 測試通過,内核爲 conda_pytorch_p310 ,實例爲一台 ml.g5.2xlarge 實例,如下圖所示。
3. 采樣過程說明
首先假設你有一個噪聲樣本(noise sample),你把這個噪聲樣本輸入到一個已經訓練好的神經網絡中。這個神經網絡已經知道精靈圖像的樣子,它接下來的主要工作是預測噪聲。請注意:這個神經網絡預測的是噪聲而不是精靈圖像,然後我們從噪聲樣本中減去預測的噪聲,來得到更像精靈圖像的輸出結果。
Source: How Diffusion Models Work, https://learn.deeplearning.ai/diffusion-models/lesson/2/intuition ,by DeepLearning.AI
由于隻是對噪聲的預測,它并不能完全消除所有噪聲,因此需要多個步驟才能獲得高質量的樣本。比如我們希望在 500 次這樣的叠代之後,能夠得到看起來非常像精靈圖像的輸出結果。
Source: How Diffusion Models Work,https://learn.deeplearning.ai/diffusion-models/lesson/2/intuition ,by DeepLearning.AI
我們先看一段僞代碼,從算法實現上高屋建瓴地看下整個邏輯結構:
Source: How Diffusion Models Work, https://learn.deeplearning.ai/diffusion-models/lesson/2/intuition ,by DeepLearning.AI
首先我們以随機采樣噪聲樣本(random noise sample)的方式,開始這段旅程。
如果你看過一些關于穿越時間旅行的電影,這整個過程很像是一段時間旅行。想像一下你有一杯墨汁,我們實際上是在用時光倒退(step backwards)的方式;它最初是完全擴散的漆黑墨汁,然後我們會一直追溯到有第一滴墨汁滴入一杯清水的那個最初時分。
然後,我們将采樣一些額外噪聲(extra noise)。爲什麽我們需要添加一些額外噪聲,這其實是一個很有趣的話題,我們會在本文的後面部分詳細探讨這個話題。
這是你實際将原始噪聲、那個樣本傳遞回神經網絡的地方,然後你會得到一些預測的噪聲。而這種預測噪聲是經過訓練的神經網絡想要從原始噪聲中減去的噪聲,以在最後得到看起來更像精靈圖像的輸出結果。
最後我們還會用到一種名爲 “DDPM” 的采樣算法,它代表降噪擴散概率模型。
4. 導入所需的庫文件
現在我們進入通過代碼解讀擴散模型的部分。首先,我們需要導入 PyTorch 和一些 PyTorch 相關的實用庫,以及導入幫助我們設計神經網絡的一些輔助函數(helper functions)。
from typing import Dict, Tuplefrom tqdm import tqdmimport torchimport torch.nn as nnimport torch.nn.functional as Ffrom torch.utils.data import DataLoaderfrom torchvision import models, transformsfrom torchvision.utils import save_image, make_gridimport matplotlib.pyplot as pltfrom matplotlib.animation import FuncAnimation, PillowWriterimport numpy as npfrom IPython.display import HTMLfrom diffusion_utilities import *
5. 神經網絡架構設計
現在我們來設置神經網絡,我們要用它來采樣。
class ContextUnet(nn.Module): def __init__(self, in_channels, n_feat=256, n_cfeat=10, height=28): # cfeat - context features super(ContextUnet, self).__init__() # number of input channels, number of intermediate feature maps and number of classes self.in_channels = in_channels self.n_feat = n_feat self.n_cfeat = n_cfeat self.h = height #assume h == w. must be divisible by 4, so 28,24,20,16... # Initialize the initial convolutional layer self.init_conv = ResidualConvBlock(in_channels, n_feat, is_res=True)# Initialize the down-sampling path of the U-Net with two levels self.down1 = UnetDown(n_feat, n_feat) # down1 #[10, 256, 8, 8] self.down2 = UnetDown(n_feat, 2 * n_feat) # down2 #[10, 256, 4, 4] # original: self.to_vec = nn.Sequential(nn.AvgPool2d(7), nn.GELU()) self.to_vec = nn.Sequential(nn.AvgPool2d((4)), nn.GELU()) # Embed the timestep and context labels with a one-layer fully connected neural network self.timeembed1 = EmbedFC(1, 2*n_feat) self.timeembed2 = EmbedFC(1, 1*n_feat) self.contextembed1 = EmbedFC(n_cfeat, 2*n_feat) self.contextembed2 = EmbedFC(n_cfeat, 1*n_feat) # Initialize the up-sampling path of the U-Net with three levels self.up0 = nn.Sequential( nn.ConvTranspose2d(2 * n_feat, 2 * n_feat, self.h//4, self.h//4), # up-sample nn.GroupNorm(8, 2 * n_feat), # normalize nn.ReLU(), ) self.up1 = UnetUp(4 * n_feat, n_feat) self.up2 = UnetUp(2 * n_feat, n_feat) # Initialize the final convolutional layers to map to the same number of channels as the input image self.out = nn.Sequential( nn.Conv2d(2 * n_feat, n_feat, 3, 1, 1), # reduce number of feature maps #in_channels, out_channels, kernel_size, stride=1, padding=0 nn.GroupNorm(8, n_feat), # normalize nn.ReLU(), nn.Conv2d(n_feat, self.in_channels, 3, 1, 1), # map to same number of channels as input ) def forward(self, x, t, c=None): """ x : (batch, n_feat, h, w) : input image t : (batch, n_cfeat) : time step c : (batch, n_classes) : context label """ # x is the input image, c is the context label, t is the timestep, context_mask says which samples to block the context on # pass the input image through the initial convolutional layer x = self.init_conv(x) # pass the result through the down-sampling path down1 = self.down1(x) #[10, 256, 8, 8] down2 = self.down2(down1) #[10, 256, 4, 4] # convert the feature maps to a vector and apply an activation hiddenvec = self.to_vec(down2) # mask out context if context_mask == 1 if c is None: c = torch.zeros(x.shape[0], self.n_cfeat).to(x) # embed context and timestep cemb1 = self.contextembed1(c).view(-1, self.n_feat * 2, 1, 1) # (batch, 2*n_feat, 1,1) temb1 = self.timeembed1(t).view(-1, self.n_feat * 2, 1, 1) cemb2 = self.contextembed2(c).view(-1, self.n_feat, 1, 1) temb2 = self.timeembed2(t).view(-1, self.n_feat, 1, 1) #print(f"uunet forward: cemb1 {cemb1.shape}. temb1 {temb1.shape}, cemb2 {cemb2.shape}. temb2 {temb2.shape}") up1 = self.up0(hiddenvec) up2 = self.up1(cemb1*up1 + temb1, down2) # add and multiply embeddings up3 = self.up2(cemb2*up2 + temb2, down1) out = self.out(torch.cat((up3, x), 1)) return out
6. 設置模型訓練的超參數
接下來,我們将設置模型訓練需要的一些超參數,包括:時間步長、圖像尺寸等。
如果對照 DDPM 的論文,其中定義了一個 noise schedule 的概念, noise schedule 決定了在特定時間裏步長對圖像施加的噪點水平。因此,這部分隻是構造一些你記得的縮放因子的 DDPM 算法參數。那些縮放值 S1、S2、S3 ,這些縮放值是在 noise schedule 中計算的。它之所以被稱爲 “Schedule”,是因爲它取決于時間步長。
-
《DDPM》
https://arxiv.org/pdf/2006.11239.pdf
Source: How Diffusion Models Work, https://learn.deeplearning.ai/diffusion-models/lesson/2/intuition ,by DeepLearning.AI
超參數介紹:
-
beta1:DDPM 算法的超參數
-
beta2:DDPM 算法的超參數
-
height:圖像的長度和高度
-
noise schedule(噪聲調度):确定在某個時間步長應用于圖像的噪聲級别;
-
S1,S2,S3:縮放因子的值
如下面代碼所示,我們在這裏設置的時間步長(timesteps)是 500 ;圖像尺寸參數 height 設置爲 16 ,表示這是 16 乘 16 的正方形圖像;DDPM 的超參數 beta1 和 beta2 等等。
# hyperparameters # diffusion hyperparameterstimesteps = 500beta1 = 1e-4beta2 = 0.02 # network hyperparametersdevice = torch.device("cuda:0" if torch.cuda.is_available() else torch.device('cpu'))n_feat = 64 # 64 hidden dimension featuren_cfeat = 5 # context vector is of size 5height = 16 # 16x16 imagesave_dir = './weights/'
請記住,你正在浏覽 500 次的步驟,因爲你正在經曆你在這裏看到的緩慢去除噪音的 500 次叠代。
Source: How Diffusion Models Work, https://learn.deeplearning.ai/diffusion-models/lesson/2/intuition ,by DeepLearning.AI
以下代碼塊将構建 DDPM 論文中定義的時間步長(noise schedule):
# construct DDPM noise scheduleb_t = (beta2 - beta1) * torch.linspace(0, 1, timesteps + 1, device=device) + beta1a_t = 1 - b_tab_t = torch.cumsum(a_t.log(), dim=0).exp() ab_t[0] = 1
接下來實例化模型:
# construct modelnn_model = ContextUnet(in_channels=3, n_feat=n_feat, n_cfeat=n_cfeat, height=height).to(device)
左滑查看更多
7. 添加額外噪聲的輸出測試
首先測試的是添加額外噪聲的輸出測試。可以重點關注下變量 z 。
在每次叠代之後,我們通過設置“z = torch.randn_like(x)”來添加額外的采樣噪聲,以讓噪聲輸入符合正态分布:
# helper function; removes the predicted noise (but adds some noise back in to avoid collapse)def denoise_add_noise(x, t, pred_noise, z=None): if z is None: z = torch.randn_like(x) noise = b_t.sqrt()[t] * z mean = (x - pred_noise * ((1 - a_t[t]) / (1 - ab_t[t]).sqrt())) / a_t[t].sqrt()
接下來加載該模型:
# load in model weights and set to eval modenn_model.load_state_dict(torch.load(f"{save_dir}/model_trained.pth", map_location=device))nn_model.eval()print("Loaded in Model")
以下代碼段實現了前面介紹過的 DDPM 采樣算法:
# sample using standard algorithm@torch.no_grad()def sample_ddpm(n_sample, save_rate=20): # x_T ~ N(0, 1), sample initial noise samples = torch.randn(n_sample, 3, height, height).to(device) # array to keep track of generated steps for plotting intermediate = [] for i in range(timesteps, 0, -1): print(f'sampling timestep {i:3d}', end='\r') # reshape time tensor t = torch.tensor([i / timesteps])[:, None, None, None].to(device) # sample some random noise to inject back in. For i = 1, don't add back in noise z = torch.randn_like(samples) if i > 1 else 0 eps = nn_model(samples, t) # predict noise e_(x_t,t) samples = denoise_add_noise(samples, i, eps, z) if i % save_rate ==0 or i==timesteps or i
-
-
-