譯者 | 朱先忠
審校 | 重樓
簡介
在我最近發表的幾篇文章中，我談到了生成式深度學習算法，這些算法大多與文本生成任務有關。所以，我認為現在轉向圖像生成的生成算法研究會很有趣。我們知道，如今已經有很多專門用于生成圖像的深度學習模型，例如自動編碼器、變分自動編碼器（VAE）、生成對抗網絡（GAN）和神經風格遷移（NST）。
在本文中，我想討論所謂的擴散模型 ——這是深度學習領域中對圖像生成影響最大的模型之一。該算法的想法最早是在2015年由Sohl-Dickstein等人撰寫的題為《使用非平衡熱力學的深度無監督學習》的論文中提出的【引文1】。他們的框架隨后由Ho等人在2020年的論文《去噪擴散概率模型》中進一步開發【引文2】。DDPM后來被OpenAI和Google改編以開發DALLE-2和Imagen，我們知道這些模型具有生成高質量圖像的令人印象深刻的能力。
擴散模型的工作原理
一般來說，擴散模型的工作原理是從噪聲中生成圖像。我們可以把它想象成一個藝術家將畫布上的一抹顏料變成一幅美麗的藝術品。為了做到這一點，首先需要訓練擴散模型。訓練模型需要遵循兩個主要步驟，即前向擴散和后向擴散。
圖1.正向和反向擴散過程【引文3】
如上圖所示，前向擴散是一個將高斯噪聲迭代地應用于原始圖像的過程。我們不斷添加噪聲，直到圖像完全無法識別，此時我們可以說圖像現在位于潛在空間中。與自動編碼器和GAN中的潛在空間通常比原始圖像的維度低不同，DDPM中的潛在空間保持與原始圖像完全相同的維度。這個噪聲過程遵循馬爾可夫鏈的原理，這意味著時間步t處的圖像僅受時間步t-1的影響。前向擴散被認為很容易，因為我們基本上只是一步一步地添加一些噪聲。
第二個訓練階段稱為反向擴散，我們的目標是一點一點地去除噪聲，直到獲得清晰的圖像。這個過程遵循逆馬爾可夫鏈的原理，其中時間步長t-1的圖像只能基于時間步長t的圖像獲得。這樣的去噪過程非常困難，因為我們需要猜測哪些像素是噪聲，哪些像素屬于實際圖像內容。因此，我們需要采用神經網絡模型來實現這一點。
DDPM使用U-Net作為反向擴散深度學習架構的基礎。但是，我們不需要使用原始的U-Net模型【引文4】，而是需要對其進行一些修改，以使其更適合我們的任務。稍后，我將在MNIST手寫數字數據集【引文5】上訓練此模型，看看它是否可以生成類似的圖像。
好了，以上就是目前你需要了解的有關擴散模型的所有基本概念。在接下來的內容中，我們將從頭開始實現擴散模型算法，從而使讀者更深入地了解有關細節。
PyTorch實現
我們將首先導入所需的模塊。如果你還不熟悉下面的導入，那么告訴你torch和torchvision都是我們用于準備模型和數據集的庫。同時，matplotlib和tqdm將幫助我們顯示圖像和進度條。
# Codeblock 1
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from tqdm import tqdm
由于模塊已導入，接下來要做的是初始化一些配置參數。詳細信息請參閱下面的Codeblock 2。
# Codeblock 2
IMAGE_SIZE = 28 #(1)
NUM_CHANNELS = 1 #(2)
BATCH_SIZE = 2
NUM_EPOCHS = 10
LEARNING_RATE = 0.001
NUM_TIMESTEPS = 1000 #(3)
BETA_START = 0.0001 #(4)
BETA_END = 0.02 #(5)
TIME_EMBED_DIM = 32 #(6)
DEVICE = torch.device("cuda" if torch.cuda.is_available else "cpu") #(7)
DEVICE
# Codeblock 2 Output
device(type='cuda')在代碼行#(1)和#(2)中，我設置IMAGE_SIZE和NUM_CHANNELS兩個常量分別為28和1，這些數字是從MNIST數據集中的圖像維度獲得的。變量BATCH_SIZE，NUM_EPOCHS和LEARNING_RATE非常簡單，所以我不需要進一步解釋它們。
在代碼行#(3)，變量NUM_TIMESTEPS表示前向和后向擴散過程中的迭代次數。時間步長0是圖像處于原始狀態的情況（圖1中最左邊的圖像）。在這種情況下，由于我們將此參數設置為1000，時間步長數999將是圖像完全無法識別的情況（上面圖1中最右邊的圖像）。重要的是要記住，時間步長的選擇涉及模型準確性和計算成本之間的權衡。如果我們為分配一個較小的值NUM_TIMESTEPS，推理時間會更短，但由于模型在后向擴散階段細化圖像的步驟較少，因此生成的圖像可能不是很好。另一方面，增加NUM_TIMESTEPS會減慢推理過程，但我們可以預期輸出圖像具有更好的質量，這要歸功于逐步的去噪過程，從而可以實現更精確的重建。
接下來，BETA_START（代碼行#(4)）和BETA_END（代碼行#(5)）變量分別用于控制每個時間步添加的高斯噪聲量，而TIME_EMBED_DIM（代碼行#(6)）用于確定用于存儲時間步信息的特征向量長度。最后，在第#(7)行，如果Pytorch檢測到我們的機器上安裝了GPU，我將分配“cuda”給變量。我強烈建議你在GPU上運行此項目，因為訓練擴散模型的計算成本很高。除了上述參數之外，為NUM_TIMESTEPS、BETA_START和BETA_END設定的值均直接采用自DDPM論文【引文2】。
完整的實現過程可分為幾個步驟：構建U-Net模型、準備數據集、定義擴散過程的噪聲調度程序、訓練和推理。我們將在以下小節中討論每個階段。
U-Net架構：時間嵌入
正如我之前提到的，擴散模型的基礎是U-Net。之所以使用這種架構，是因為它的輸出層適合表示圖像，這絕對是有道理的，因為它最初是為圖像分割任務引入的。下圖顯示了原始U-Net架構的樣子。
圖2.【引文4】中提出的原始U-Net模型
然而，有必要修改一下此架構，以便它也能考慮時間步長信息。不僅如此，由于我們只使用MNIST數據集，我們還需要使模型更小。只需記住深度學習中的慣例，即更簡單的模型通常對簡單任務更有效。
下圖中我展示了經過修改的整個U-Net模型。在這里你可以看到時間嵌入張量在每個階段都被注入到模型中，稍后將通過逐元素求和來完成，從而使模型能夠捕獲時間步長信息。接下來，我們不會像原始U-Net那樣將每個下采樣和上采樣階段重復四次，而是在本例中只重復兩次。此外，值得注意的是，下采樣階段的堆棧也稱為編碼器，而上采樣階段的堆棧通常稱為解碼器。
圖3.針對我們的擴散任務修改后的U-Net模型【引文3】
現在，讓我們開始構建架構，創建一個用于生成時間嵌入張量的類，其思想類似于Transformer中的位置嵌入。有關詳細信息，請參閱下面的Codeblock 3。
# Codeblock 3
class TimeEmbedding(nn.Module):
def forward(self):
time = torch.arange(NUM_TIMESTEPS, device=DEVICE).reshape(NUM_TIMESTEPS, 1) #(1)
print(f"time\t\t: {time.shape}")
i = torch.arange(0, TIME_EMBED_DIM, 2, device=DEVICE)
denominator = torch.pow(10000, i/TIME_EMBED_DIM)
print(f"denominator\t: {denominator.shape}")
even_time_embed = torch.sin(time/denominator) #(1)
odd_time_embed = torch.cos(time/denominator) #(2)
print(f"even_time_embed\t: {even_time_embed.shape}")
print(f"odd_time_embed\t: {odd_time_embed.shape}")
stacked = torch.stack([even_time_embed, odd_time_embed], dim=2) #(3)
print(f"stacked\t\t: {stacked.shape}")
time_embed = torch.flatten(stacked, start_dim=1, end_dim=2) #(4)
print(f"time_embed\t: {time_embed.shape}")
return time_embed上述代碼中，我們基本上要創建一個大小為NUM_TIMESTEPS×TIME_EMBED_DIM（1000×32）的張量，其中該張量的每一行都包含時間步長信息。稍后，1000個時間步長中的每一個都將由長度為32的特征向量表示。張量本身的值是根據圖4中的兩個方程獲得的。在上面的Codeblock 3中，這兩個方程分別在第#(1)和#(2)行實現，每個方程都形成一個大小為1000×16的張量。接下來，使用第#(3)行和#(4)行的代碼將這些張量組合起來。
在這里，我還打印出了上述代碼塊中完成的每個步驟，以便你更好地理解TimeEmbedding類中實際執行的操作。如果你仍想了解有關上述代碼的更多解釋，請隨時閱讀我之前關于Transformer的博客文章，你可以通過本文末尾的鏈接訪問該文章。單擊鏈接后，你可以一直向下滾動到“Positional Encoding”部分。
圖4. Transformer論文【引文6】中的正弦位置編碼公式
現在，讓我們使用以下測試代碼檢查是否TimeEmbedding類正常工作。結果輸出顯示它成功生成了一個大小為1000×32的張量，這正是我們之前所期望的。
# Codeblock 4
time_embed_test = TimeEmbedding()
out_test = time_embed_test()
# Codeblock 4 Output
time : torch.Size([1000, 1])
denominator : torch.Size([16])
even_time_embed : torch.Size([1000, 16])
odd_time_embed : torch.Size([1000, 16])
stacked : torch.Size([1000, 16, 2])
time_embed : torch.Size([1000, 32])U-Net架構：DoubleConv
如果仔細觀察修改后的架構，你會發現我們實際上得到了許多重復的模型，例如下圖中黃色框中突出顯示的模型。
圖5.黃色框內完成的流程將在DoubleConv課堂上實現【引文3】
這五個黃色框具有相同的結構，它們由兩個卷積層組成，在執行第一個卷積操作后立即注入時間嵌入張量。因此，我們現在要做的是創建另一個名為DoubleConv的類來重現此結構。查看下面的代碼塊5a和5b以了解我如何做到這一點。
# Codeblock 5a
class DoubleConv(nn.Module):
def __init__(self, in_channels, out_channels): #(1)
super().__init__()
self.conv_0 = nn.Conv2d(in_channels=in_channels, #(2)
out_channels=out_channels,
kernel_size=3,
bias=False,
padding=1)
self.bn_0 = nn.BatchNorm2d(num_features=out_channels) #(3)
self.time_embedding = TimeEmbedding() #(4)
self.linear = nn.Linear(in_features=TIME_EMBED_DIM, #(5)
out_features=out_channels)
self.conv_1 = nn.Conv2d(in_channels=out_channels, #(6)
out_channels=out_channels,
kernel_size=3,
bias=False,
padding=1)
self.bn_1 = nn.BatchNorm2d(num_features=out_channels) #(7)
self.relu = nn.ReLU(inplace=True) #(8)上述方法__init__()中的兩個輸入參數使我們可以靈活地配置輸入和輸出通道的數量（#(1)）；這樣一來，DoubleConv類可以通過調整輸入參數來實例化所有五個黃色框。顧名思義，這里我們初始化了兩個卷積層（行#(2)和#(6)），每個卷積層后跟一個批量歸一化層和一個ReLU激活函數。請記住，兩個歸一化層需要單獨初始化（行#(3)和#(7)），因為它們每個都有自己可訓練的歸一化參數。同時，ReLU激活函數應該只初始化一次（#(8)），因為它不包含任何參數，因此可以在網絡的不同部分多次使用。在代碼行#(4)，我們初始化之前創建的TimeEmbedding層，該層稍后將連接到標準線性層（#(5)）。該線性層負責調整時間嵌入張量的維度，以便可以以逐元素的方式將得到的輸出與第一個卷積層的輸出相加。
現在，讓我們看一下下面的Codeblock 5b，以更好地理解該DoubleConv塊的運行流程。在這部分代碼中，你可以看到forward()方法接受兩個輸入：原始圖像x和時間步長信息t，如第#(1)行所示。我們首先使用第一個Conv-BN-ReLU序列處理圖像（#(2–4)）。這種Conv-BN-ReLU結構通常用于基于CNN的模型，即使插圖沒有明確顯示批量標準化和ReLU層。除了圖像之外，我們還從相應圖像的嵌入張量（#(5)）中獲取第t個時間步長信息并將其傳遞給線性層（#(6)）。我們仍然需要使用第#(7)行的代碼擴展結果張量的維度，然后在第#(8)行執行逐元素求和。最后，我們使用第二個Conv-BN-ReLU序列(#(9–11))處理結果張量。
# Codeblock 5b
def forward(self, x, t): #(1)
print(f'images\t\t\t: {x.size()}')
print(f'timesteps\t\t: {t.size()}, {t}')
x = self.conv_0(x) #(2)
x = self.bn_0(x) #(3)
x = self.relu(x) #(4)
print(f'\nafter first conv\t: {x.size()}')
time_embed = self.time_embedding()[t] #(5)
print(f'\ntime_embed\t\t: {time_embed.size()}')
time_embed = self.linear(time_embed) #(6)
print(f'time_embed after linear\t: {time_embed.size()}')
time_embed = time_embed[:, :, None, None] #(7)
print(f'time_embed expanded\t: {time_embed.size()}')
x = x + time_embed #(8)
print(f'\nafter summation\t\t: {x.size()}')
x = self.conv_1(x) #(9)
x = self.bn_1(x) #(10)
x = self.relu(x) #(11)
print(f'after second conv\t: {x.size()}')
return x為了查看我們的DoubleConv類實現是否正常工作，我們將使用下面的Codeblock 6對其進行測試。在這里，我想模擬此塊的第一個實例，它對應于圖5中最左邊的黃色框。為此，我們需要將in_channels和out_channels參數分別設置為1和64（#(1)）。接下來，我們初始化兩個輸入張量，即x_test和t_test。其中，x_test張量的大小為2×1×28×28，表示一批大小為28×28的兩個灰度圖像（#(2)）。請記住，這只是一個賦予了隨機數值的虛擬張量，它將在訓練階段的后期被來自MNIST數據集的實際圖像替換。同時，t_test是一個包含相應圖像的時間步長數的張量（#(3)）。此張量的值在0到NUM_TIMESTEPS（1000）之間隨機選擇。請注意，此張量的數據類型必須是整數，因為這些數字將用于索引，如Codeblock 5b中的第#(5)行所示。最后，在代碼行#(4)我們將x_test和t_test張量傳遞給double_conv_test層。
順便說一句，print()在運行以下代碼之前，我重新運行了前面的代碼塊并刪除了函數，以便輸出結果看起來更整潔一些。
# Codeblock 6
double_conv_test = DoubleConv(in_channels=1, out_channels=64).to(DEVICE) #(1)
x_test = torch.randn((BATCH_SIZE, NUM_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)).to(DEVICE) #(2)
t_test = torch.randint(0, NUM_TIMESTEPS, (BATCH_SIZE,)).to(DEVICE) #(3)
out_test = double_conv_test(x_test, t_test) #(4)
# Codeblock 6 Output
images : torch.Size([2, 1, 28, 28]) #(1)
timesteps : torch.Size([2]), tensor([468, 304], device='cuda:0') #(2)
after first conv : torch.Size([2, 64, 28, 28]) #(3)
time_embed : torch.Size([2, 32]) #(4)
time_embed after linear : torch.Size([2, 64])
time_embed expanded : torch.Size([2, 64, 1, 1]) #(5)
after summation : torch.Size([2, 64, 28, 28]) #(6)
after second conv : torch.Size([2, 64, 28, 28]) #(7)原始輸入張量的形狀可以在上面的輸出的第#(1)和#(2)行看到。具體來說，在第#(2)行，我還打印出了我們隨機選擇的兩個時間步。在這個例子中，我們假設x張量中的兩個圖像在輸入網絡之前已經用來自第468和第304個時間步的噪聲級別進行了噪聲處理。我們可以看到，在經過第一個卷積層（#(3)行）后，圖像張量x的形狀變為2×64×28×28。同時，我們的時間嵌入張量的大小變為2×32（#(4)行），這是通過從大小為1000×32的原始嵌入中提取第468行和第304行獲得的。為了能夠執行元素級求和（#(6)行），我們需要將32維時間嵌入向量映射到64維并擴展它們的軸，從而得到大小為2×64×1×1的張量（#(5)行），以便可以將其廣播到2×64×28×28張量。求和完成后，我們將張量傳遞到第二個卷積層，此時張量維度完全不會改變（#(7)行）。
U-Net架構：編碼器
到目前為止，既然我們已經成功實現了DoubleConv塊，那么接下來要做的就是實現所謂的DownSample塊。在下面的圖6中，這對應于紅色框內的部分。
圖6.網絡中以紅色突出顯示的部分是所謂的DownSample塊【引文3】
DownSample塊的目的是減少圖像的空間維度，但需要注意的是，它同時增加了通道數。為了實現這一點，我們可以簡單地堆疊一個DoubleConv塊和一個最大池化操作。在這種情況下，池化使用2×2內核大小，步長為2，導致圖像的空間維度是輸入的兩倍。此塊的實現可以在下面的Codeblock 7中看到。
# Codeblock 7
class DownSample(nn.Module):
def __init__(self, in_channels, out_channels): #(1)
super().__init__()
self.double_conv = DoubleConv(in_channels=in_channels, #(2)
out_channels=out_channels)
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2) #(3)
def forward(self, x, t): #(4)
print(f'original\t\t: {x.size()}')
print(f'timesteps\t\t: {t.size()}, {t}')
convolved = self.double_conv(x, t) #(5)
print(f'\nafter double conv\t: {convolved.size()}')
maxpooled = self.maxpool(convolved) #(6)
print(f'after pooling\t\t: {maxpooled.size()}')
return convolved, maxpooled #(7)在這里，我將__init__()方法設置為采用輸入和輸出通道的數量，以便我們可以使用它來創建圖6中突出顯示的兩個DownSample塊，而無需將它們寫入單獨的類（#(1)）。接下來，DoubleConv和最大池化層分別在第#(2)和#(3)行初始化。請記住，由于DoubleConv塊接受圖像x和相應的時間步t作為輸入，我們還需要設置此DownSample塊的forward()方法，以便使其也接受它們兩個參數（#(4)）。然后，當double_conv層處理兩個張量時，x和t中包含的信息被組合在一起，輸出被存儲在名為convolved（#（5））的變量中。之后，我們現在實際上在第1行使用最大池化操作執行下采樣（#(6)），產生一個名為maxpooled的張量。值得注意的是，convolved和maxpooled張量都將被返回，這基本上是因為我們稍后會把maxpooled帶入下一個下采樣階段，而convolved張量將通過跳過連接直接傳輸到解碼器中的上采樣階段。
現在，我們使用下面的Codeblock 8來測試該DownSample類。這里使用的輸入張量與Codeblock 6中的完全相同。根據結果輸出，我們可以看到池化操作成功地將DoubleConv塊的輸出從2×64×28×28（#(1)）轉換為2×64×14×14（#(2)），這表明我們的DownSample類正常工作。
# Codeblock 8
down_sample_test = DownSample(in_channels=1, out_channels=64).to(DEVICE)
x_test = torch.randn((BATCH_SIZE, NUM_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)).to(DEVICE)
t_test = torch.randint(0, NUM_TIMESTEPS, (BATCH_SIZE,)).to(DEVICE)
out_test = down_sample_test(x_test, t_test)
# Codeblock 8 Output
original : torch.Size([2, 1, 28, 28])
timesteps : torch.Size([2]), tensor([468, 304], device='cuda:0')
after double conv : torch.Size([2, 64, 28, 28]) #(1)
after pooling : torch.Size([2, 64, 14, 14]) #(2)U-Net架構：解碼器
我們需要在解碼器中引入所謂的UpSample塊，它負責將中間層中的張量恢復到原始圖像維度。為了保持對稱結構，UpSample塊的數量必須與DownSample塊的數量相匹配。觀察下面的圖7，就可以看到兩個UpSample塊的位置。
圖7.藍色框內的組件就是所謂的UpSample塊【引文3】
由于兩個UpSample塊的結構相同，我們可以為它們初始化一個類，就像DownSample我們之前創建的類一樣。請查看下面的Codeblock 9，了解我如何實現它。
# Codeblock 9
class UpSample(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.conv_transpose = nn.ConvTranspose2d(in_channels=in_channels, #(1)
out_channels=out_channels,
kernel_size=2, stride=2) #(2)
self.double_conv = DoubleConv(in_channels=in_channels, #(3)
out_channels=out_channels)
def forward(self, x, t, connection): #(4)
print(f'original\t\t: {x.size()}')
print(f'timesteps\t\t: {t.size()}, {t}')
print(f'connection\t\t: {connection.size()}')
x = self.conv_transpose(x) #(5)
print(f'\nafter conv transpose\t: {x.size()}')
x = torch.cat([x, connection], dim=1) #(6)
print(f'after concat\t\t: {x.size()}')
x = self.double_conv(x, t) #(7)
print(f'after double conv\t: {x.size()}')
return x在該__init__()方法中，我們使用nn.ConvTranspose2d對空間維度進行上采樣（#(1)）。內核大小和步長都設置為2，這樣輸出將是原來的兩倍（#(2)）。接下來，DoubleConv將使用塊來減少通道數，同時結合來自時間嵌入張量的時間步長信息（#(3)）。
這個UpSample類的流程比DownSample類稍微復雜一些。如果我們仔細看看這個架構，我們會發現我們還有一個直接來自編碼器的跳過連接。因此，除了原始圖像x和時間步長之外，我們還需要forward()方法接受另一個參數t，即殘差張量connection（#(4)）。我們在這個方法中做的第一件事是用轉置卷積層處理原始圖像（#(5)）。事實上，這個層不僅對空間大小進行了上采樣，而且還同時減少了通道數。然而，得到的張量隨后以通道方式直接連接（#(6)），使其看起來好像沒有執行任何通道減少。重要的是要知道，此時這兩個張量只是連接在一起，這意味著來自兩者的信息尚未結合。我們最終將這些連接的張量輸入到double_conv層（#(7)），允許它們通過卷積層內的可學習參數相互共享信息。
下面的Codeblock 10展示了我如何測試該類UpSample。需要傳入的張量的大小是根據第二個上采樣塊（即圖7中最右邊的藍色框）設置的。
# Codeblock 10
up_sample_test = UpSample(in_channels=128, out_channels=64).to(DEVICE)
x_test = torch.randn((BATCH_SIZE, 128, 14, 14)).to(DEVICE)
t_test = torch.randint(0, NUM_TIMESTEPS, (BATCH_SIZE,)).to(DEVICE)
connection_test = torch.randn((BATCH_SIZE, 64, 28, 28)).to(DEVICE)
out_test = up_sample_test(x_test, t_test, connection_test)在下面的結果輸出中，如果我們將輸入張量（#(1)）與最終張量形狀（#(2)）進行比較，我們可以清楚地看到通道數成功地從128減少到64，同時空間維度從14×14增加到28×28。這實際上意味著，我們的UpSample類現在可以在主U-Net架構中使用了。
original : torch.Size([2, 128, 14, 14]) #(1)
timesteps : torch.Size([2]), tensor([468, 304], device='cuda:0')
connection : torch.Size([2, 64, 28, 28])
after conv transpose : torch.Size([2, 64, 28, 28])
after concat : torch.Size([2, 128, 28, 28])
after double conv : torch.Size([2, 64, 28, 28]) #(2)U-Net架構：將所有組件整合在一起
一旦創建了所有U-Net組件，我們接下來要做的就是將它們包裝成一個類。請參閱下面的代碼塊11a和11b了解詳細信息。
# Codeblock 11a
class UNet(nn.Module):
def __init__(self):
super().__init__()
self.downsample_0 = DownSample(in_channels=NUM_CHANNELS, #(1)
out_channels=64)
self.downsample_1 = DownSample(in_channels=64, #(2)
out_channels=128)
self.bottleneck = DoubleConv(in_channels=128, #(3)
out_channels=256)
self.upsample_0 = UpSample(in_channels=256, #(4)
out_channels=128)
self.upsample_1 = UpSample(in_channels=128, #(5)
out_channels=64)
self.output = nn.Conv2d(in_channels=64, #(6)
out_channels=NUM_CHANNELS,
kernel_size=1)你可以在上面的__init__()方法中看到，我們初始化兩個下采樣（#(1-2)）和兩個上采樣（#(4-5)）塊；其中，輸入和輸出通道的數量根據圖中所示的架構設置。實際上，還有兩個我還沒有解釋的額外組件，即瓶頸層bottleneck（#(3)）和輸出層output（#(6)）。前者本質上只是一個DoubleConv塊，它充當編碼器和解碼器之間的主要連接。查看下面的圖8，以查看網絡的哪些組件屬于瓶頸層bottleneck。接下來，輸出層是一個標準卷積層，負責將上一個UpSampling階段生成的64通道圖像轉換為僅1通道。此操作使用大小為1×1的內核完成，這意味著，它結合所有通道的信息，同時在每個像素位置獨立操作。
圖8.瓶頸層（模型的下部）充當U-Net編碼器和解碼器之間的主要橋梁【引文3】
我想，以下代碼塊中的整個U-Net的forward()方法應該是非常簡單的，因為我們在這里所做的基本上是將張量從一層傳遞到另一層——只是不要忘記在下采樣和上采樣塊之間包含跳過連接。
# Codeblock 11b
def forward(self, x, t): #(1)
print(f'original\t\t: {x.size()}')
print(f'timesteps\t\t: {t.size()}, {t}')
convolved_0, maxpooled_0 = self.downsample_0(x, t)
print(f'\nmaxpooled_0\t\t: {maxpooled_0.size()}')
convolved_1, maxpooled_1 = self.downsample_1(maxpooled_0, t)
print(f'maxpooled_1\t\t: {maxpooled_1.size()}')
x = self.bottleneck(maxpooled_1, t)
print(f'after bottleneck\t: {x.size()}')
upsampled_0 = self.upsample_0(x, t, convolved_1)
print(f'upsampled_0\t\t: {upsampled_0.size()}')
upsampled_1 = self.upsample_1(upsampled_0, t, convolved_0)
print(f'upsampled_1\t\t: {upsampled_1.size()}')
x = self.output(upsampled_1)
print(f'final output\t\t: {x.size()}')
return x現在，讓我們通過運行以下測試代碼來看看我們是否正確構建了上面的U-Net類。
# Codeblock 12
unet_test = UNet().to(DEVICE)
x_test = torch.randn((BATCH_SIZE, NUM_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)).to(DEVICE)
t_test = torch.randint(0, NUM_TIMESTEPS, (BATCH_SIZE,)).to(DEVICE)
out_test = unet_test(x_test, t_test)
# Codeblock 12 Output
original : torch.Size([2, 1, 28, 28]) #(1)
timesteps : torch.Size([2]), tensor([468, 304], device='cuda:0')
maxpooled_0 : torch.Size([2, 64, 14, 14]) #(2)
maxpooled_1 : torch.Size([2, 128, 7, 7]) #(3)
after bottleneck : torch.Size([2, 256, 7, 7]) #(4)
upsampled_0 : torch.Size([2, 128, 14, 14])
upsampled_1 : torch.Size([2, 64, 28, 28])
final output : torch.Size([2, 1, 28, 28]) #(5)在上面的輸出中我們可以看到，兩個下采樣階段成功地將原始大小為1×28×28(#(1))的張量分別轉換為64×14×14（#(2)）和128×7×7（#(3)）。然后，該張量通過瓶頸層，導致其通道數擴展到256，而不會改變空間維度（#(4)）。最后，我們對張量進行兩次上采樣，最終將通道數縮小到1（#(5)）。根據這個輸出，我們的模型似乎已經可以運行正常了。因此，它現在可以用于我們的擴散任務了。
數據集準備
由于我們已經成功創建了整個U-Net架構，接下來要做的就是準備MNIST手寫數字數據集。在實際加載它之前，我們需要首先使用Torchvision中的transforms.Compose()方法定義預處理步驟，如Codeblock13中#(1)行所示。這里我們做兩件事：將圖像轉換為PyTorch張量，這也會將像素值從0-255縮放到0-1（#(2)），并對其進行規范化，以便最終像素值介于-1和1之間（#(3)）。
接下來，我們使用datasets.MNIST()下載數據集。在本例中，我們將從訓練數據中獲取圖像，因此我們使用train=True（#(5)）。不要忘記將我們之前初始化的變量transform傳遞給transform參數（transform=transform），以便它在加載圖像時自動預處理圖像（#(6)）。最后，我們需要使用DataLoader加載來自mnist_dataset的圖像（#(7)）。注意，我用于輸入參數的參數，為的是在每次迭代中從數據集中隨機挑選BATCH_SIZE（2）張圖像。
# Codeblock 13
transform = transforms.Compose([ #(1)
transforms.ToTensor(), #(2)
transforms.Normalize((0.5,), (0.5,)) #(3)
])
mnist_dataset = datasets.MNIST( #(4)
root='./data',
train=True, #(5)
download=True,
transform=transform #(6)
)
loader = DataLoader(mnist_dataset, #(7)
batch_size=BATCH_SIZE,
drop_last=True,
shuffle=True)在下面的代碼塊中，我嘗試從數據集中加載一批圖像。在每次迭代中，loader都會提供圖像和相應的標簽，因此我們需要將它們存儲在兩個單獨的變量中：images和labels。
# Codeblock 14
images, labels = next(iter(loader))
print('images\t\t:', images.shape)
print('labels\t\t:', labels.shape)
print('min value\t:', images.min())
print('max value\t:', images.max())我們可以在下面的結果輸出中看到，images張量的大小為2×1×28×28（#(1)），表示已成功加載兩張大小為28×28的灰度圖像。在這里我們還可以看到labels張量的長度為2，與加載的圖像數量相匹配（#(2)）。請注意，在這種情況下，標簽將被完全忽略。我的計劃是，我只希望模型從整個訓練數據集中生成它之前看到的任何數字，甚至不知道它實際上是什么數字。最后，此輸出還顯示預處理工作正常，因為像素值現在介于-1和1之間。
# Codeblock 14 Output
images : torch.Size([2, 1, 28, 28]) #(1)
labels : torch.Size([2]) #(2)
min value : tensor(-1.)
max value : tensor(1.)
如果你想看看我們剛剛加載的圖像是什么樣子，請運行以下代碼。
# Codeblock 15
plt.imshow(images[0].squeeze(), cmap='gray')
plt.show()
圖9.Codeblock 15的輸出【引文3】
噪音調度器
在本節中，我們將討論如何執行前向和后向擴散，該過程本質上涉及在每個時間步長上一點一點地添加或消除噪聲。有必要知道，我們基本上希望在所有時間步長上噪聲量均勻；其中，在前向擴散中，圖像應該在時間步長1000時完全充滿噪聲，而在后向擴散中，我們必須在時間步長0時獲得完全清晰的圖像。因此，我們需要一些東西來控制每個時間步長的噪聲量。在本節后面，我將實現一個名為NoiseScheduler的類來執行此操作。這可能是本文中數學知識涉及最多的部分，因為我將在這里展示許多方程式。但不要擔心，因為我們將專注于實現這些方程式，而不是討論數學推導。
現在，讓我們看一下圖10中的方程式，我將在下面的NoiseScheduler類的__init__()方法中實現它們。
圖10.我們需要在類NoiseScheduler的__init__()方法中實現的方程【引文3】
# Codeblock 16a
class NoiseScheduler:
def __init__(self):
self.betas = torch.linspace(BETA_START, BETA_END, NUM_TIMESTEPS) #(1)
self.alphas = 1. - self.betas
self.alphas_cum_prod = torch.cumprod(self.alphas, dim=0)
self.sqrt_alphas_cum_prod = torch.sqrt(self.alphas_cum_prod)
self.sqrt_one_minus_alphas_cum_prod = torch.sqrt(1. - self.alphas_cum_prod)上述代碼通過創建多個數字序列來工作，它們基本上都由BETA_START（0.0001）、BETA_END（0.02）和NUM_TIMESTEPS（1000）控制。我們需要實例化的第一個序列是betas本身，它是使用torch.linspace()完成的（#(1)）。它本質上的作用是生成一個長度為1000的一維張量，從0.0001到0.02，其中此張量中的每個元素都對應一個時間步。每個元素之間的間隔是均勻的，這使我們能夠在所有時間步長中生成均勻的噪聲量。有了這個betas張量，我們然后根據圖10中的四個方程計算alphas、alphas_cum_prod、sqrt_alphas_cum_prod和sqrt_one_minus_alphas_cum_prod。稍后，這些張量將作為在擴散過程中如何產生或消除噪聲的基礎。
擴散通常以順序方式進行。然而，前向擴散過程是確定性的，因此我們可以將原始方程推導為封閉形式，這樣我們就可以在特定的時間步長中獲得噪聲，而不必從頭開始迭代添加噪聲。下圖11顯示了前向擴散的封閉形式，其中x0表示原始圖像，而epsilon（?）表示由隨機高斯噪聲組成的圖像。我們可以將此方程視為加權組合，其中我們根據時間步長確定的權重將清晰圖像和噪聲組合在一起，從而得到具有特定噪聲量的圖像。
圖11.正向擴散過程的封閉形式【引文3】
該方程的實現可以在Codeblock 16b中看到。在forward_diffusion()方法中，x0和?表示為original和noise。這里你需要記住這兩個輸入變量是圖像，而sqrt_alphas_cum_prod_t和sqrt_one_minus_alphas_cum_prod_t是標量。因此，我們需要調整這兩個標量的形狀（代碼行#(1)和#(2)），以便可以執行#(3)行中的操作。變量noisy_image將是此函數的輸出，我想這個名字是不言自明的。
# Codeblock 16b
def forward_diffusion(self, original, noise, t):
sqrt_alphas_cum_prod_t = self.sqrt_alphas_cum_prod[t]
sqrt_alphas_cum_prod_t = sqrt_alphas_cum_prod_t.to(DEVICE).view(-1, 1, 1, 1) #(1)
sqrt_one_minus_alphas_cum_prod_t = self.sqrt_one_minus_alphas_cum_prod[t]
sqrt_one_minus_alphas_cum_prod_t = sqrt_one_minus_alphas_cum_prod_t.to(DEVICE).view(-1, 1, 1, 1) #(2)
noisy_image = sqrt_alphas_cum_prod_t * original + sqrt_one_minus_alphas_cum_prod_t * noise #(3)
return noisy_image現在，我們來談談反向擴散。事實上，這個比正向擴散稍微復雜一些，因為我們需要三個方程。在我向你展示這些方程之前，讓我先向你展示一下有關代碼實現。請參閱下面的代碼塊16c。
# Codeblock 16c
def backward_diffusion(self, current_image, predicted_noise, t): #(1)
denoised_image = (current_image - (self.sqrt_one_minus_alphas_cum_prod[t] * predicted_noise)) / self.sqrt_alphas_cum_prod[t] #(2)
denoised_image = 2 * (denoised_image - denoised_image.min()) / (denoised_image.max() - denoised_image.min()) - 1 #(3)
current_prediction = current_image - ((self.betas[t] * predicted_noise) / (self.sqrt_one_minus_alphas_cum_prod[t])) #(4)
current_prediction = current_prediction / torch.sqrt(self.alphas[t]) #(5)
if t == 0: #(6)
return current_prediction, denoised_image
else:
variance = (1 - self.alphas_cum_prod[t-1]) / (1. - self.alphas_cum_prod[t]) #(7)
variance = variance * self.betas[t] #(8)
sigma = variance ** 0.5
z = torch.randn(current_image.shape).to(DEVICE)
current_prediction = current_prediction + sigma*z
return current_prediction, denoised_image在推理階段的后期，該backward_diffusion()方法將在一個循環中調用，該循環迭代NUM_TIMESTEPS(1000)次，從t=999開始，繼續執行t=998，一直到t=0為止。此函數負責根據current_image（前一個去噪步驟生成的圖像）、predicted_noise（U-Net在上一步中預測的噪聲）和時間步長信息t三個參數迭代地從圖像中去除噪聲（#(1)）。在每次迭代中，使用圖12中所示的公式進行噪聲消除，在代碼塊16c中，這對應于代碼行#(4-5)。
圖12.用于從圖像中去除噪聲的方程【引文3】
只要我們還沒有達到t=0，我們就會根據圖13中的方程計算方差（代碼行#(7-8)）。然后，該方差將用于引入另一個受控噪聲，以模擬反向擴散過程中的隨機性，因為圖12中的噪聲消除方程是確定性近似。這本質上也是我們在達到t=0（代碼行#(6)）后不再計算方差的原因，因為我們不再需要添加更多噪聲，因為圖像已經完全清晰了。
圖13.用于計算引入受控噪聲的方差的方程【引文3】
current_prediction與旨在估計前一個時間步長的圖像（xt-1）不同，張量的目標denoised_image是重建原始圖像（x0）。由于這些不同的目標，我們需要一個單獨的方程來計算denoised_image，如下圖14所示。方程本身的實現寫在第#(2-3)行。
圖14.重建原始圖像的方程【引文3】
現在，讓我們測試一下上面創建的NoiseScheduler類。在下面的代碼塊中，我實例化一個NoiseScheduler對象并打印出與其關聯的屬性，這些屬性都是使用圖10中的公式根據存儲在betas屬性中的值計算得出的。請記住，這些張量的實際長度是NUM_TIMESTEPS（1000），但在這里我只打印出前6個元素。
# Codeblock 17
noise_scheduler = NoiseScheduler()
print(f'betas\t\t\t\t: {noise_scheduler.betas[:6]}')
print(f'alphas\t\t\t\t: {noise_scheduler.alphas[:6]}')
print(f'alphas_cum_prod\t\t\t: {noise_scheduler.alphas_cum_prod[:6]}')
print(f'sqrt_alphas_cum_prod\t\t: {noise_scheduler.sqrt_alphas_cum_prod[:6]}')
print(f'sqrt_one_minus_alphas_cum_prod\t: {noise_scheduler.sqrt_one_minus_alphas_cum_prod[:6]}')
# Codeblock 17 Output
betas : tensor([1.0000e-04, 1.1992e-04, 1.3984e-04, 1.5976e-04, 1.7968e-04, 1.9960e-04])
alphas : tensor([0.9999, 0.9999, 0.9999, 0.9998, 0.9998, 0.9998])
alphas_cum_prod : tensor([0.9999, 0.9998, 0.9996, 0.9995, 0.9993, 0.9991])
sqrt_alphas_cum_prod : tensor([0.9999, 0.9999, 0.9998, 0.9997, 0.9997, 0.9996])
sqrt_one_minus_alphas_cum_prod : tensor([0.0100, 0.0148, 0.0190, 0.0228, 0.0264, 0.0300])上面的輸出表明我們的__init__()方法按預期工作。接下來，我們將測試forward_diffusion()方法。如果回看一下圖16b，你將看到forward_diffusion()接受三個輸入：原始圖像、噪聲圖像和時間步長數。我們只需使用我們之前加載的MNIST數據集中的圖像作為第一個輸入（#(1)），并使用大小完全相同的隨機高斯噪聲作為第二個輸入（#(2)）。為此，可以運行下面的Codeblock 18來查看這兩個圖像是什么樣子。
# Codeblock 18
image = images[0] #(1)
noise = torch.randn_like(image) #(2)
plt.imshow(image.squeeze(), cmap='gray')
plt.show()
plt.imshow(noise.squeeze(), cmap='gray')
plt.show()
圖15.用作原始圖像（左）和噪聲圖像（右）的兩幅圖像（其中，左側的圖像與我之前在圖9【引文3】中展示的圖像相同）
由于我們已經準備好了圖像和噪聲，接下來我們需要做的就是將它們傳遞給forward_diffusion()方法。我實際上嘗試多次運行下面的Codeblock19：t=50、100、150等等，直到t=300。你可以在圖16中看到，隨著參數的增加，圖像變得不那么清晰。在這種情況下，當t設置為999時，圖像將完全被噪聲填充。
# Codeblock 19
noisy_image_test = noise_scheduler.forward_diffusion(image.to(DEVICE), noise.to(DEVICE), t=50)
plt.imshow(noisy_image_test[0].squeeze().cpu(), cmap='gray')
plt.show()
圖16. t=50、100、150等時刻正向擴散過程的結果，直到t=300【引文3】
不幸的是，我們無法測試該backward_diffusion()方法，因為此過程需要我們對U-Net模型進行訓練。所以，我們現在就跳過這部分。我將向你展示如何在稍后的推理階段實際使用此功能。
訓練
由于U-Net模型、MNIST數據集和噪聲調度程序已準備就緒，我們現在可以準備一個函數進行訓練。在此之前，我在下面的Codeblock 20中實例化了模型和噪聲調度程序。
# Codeblock 20
model = UNet().to(DEVICE)
noise_scheduler = NoiseScheduler()整個訓練過程在Codeblock 21中顯示的train()函數中實現。在執行任何操作之前，我們首先初始化優化器和損失函數，在本例中我們分別使用Adam和MSE（#(1-2)）。我們基本上想要做的是訓練模型，使其能夠預測輸入圖像中包含的噪聲，稍后，預測的噪聲將用作后向擴散階段去噪過程的基礎。要實際訓練模型，我們首先需要使用#(6)行處的代碼執行前向擴散。此噪聲化過程將使用#(4)行處生成的隨機噪聲在images量上完成（#(3)）。接下來，我們為t（#(5)）取0到NUM_TIMESTEPS（1000）之間的隨機數，這主要是因為我們希望我們的模型能夠看到不同噪聲水平的圖像，以此作為提高泛化能力的一種方法。生成噪聲圖像后，我們將它與所選的t（#(7)）一起傳遞給U-Net模型。此處的輸入對模型很有用，因為它表示圖像中的當前噪聲水平。最后，我們之前初始化的損失函數負責計算實際噪聲與原始圖像的預測噪聲之間的差異（#(8)）。因此，這次訓練的目標基本上是使預測噪聲盡可能與我們在第#(4)行生成的噪聲相似。
# Codeblock 21
def train():
optimizer = Adam(model.parameters(), lr=LEARNING_RATE) #(1)
loss_function = nn.MSELoss() #(2)
losses = []
for epoch in range(NUM_EPOCHS):
print(f'Epoch no {epoch}')
for images, _ in tqdm(loader):
optimizer.zero_grad()
images = images.float().to(DEVICE) #(3)
noise = torch.randn_like(images) #(4)
t = torch.randint(0, NUM_TIMESTEPS, (BATCH_SIZE,)) #(5)
noisy_images = noise_scheduler.forward_diffusion(images, noise, t).to(DEVICE) #(6)
predicted_noise = model(noisy_images, t) #(7)
loss = loss_function(predicted_noise, noise) #(8)
losses.append(loss.item())
loss.backward()
optimizer.step()
return losses現在，讓我們使用下面的代碼塊運行上述訓練函數。坐下來放松一下，等待訓練完成。就我而言，我使用了開啟了Nvidia GPU P100的Kaggle Notebook，大約需要45分鐘才能完成。
# Codeblock 22
losses = train()如果我們看一下損失圖，似乎我們的模型學習得很好，因為價值通常會隨著時間的推移而下降，早期階段下降迅速，后期階段趨勢更穩定（但仍在下降）。所以，我認為我們可以在推理階段的后期期待好的結果。
# Codeblock 23
plt.plot(losses)
圖17.損失值如何隨著訓練的進行而減少【引文3】
推理
此時，我們已經訓練了模型，因此我們現在可以對其進行推理。查看下面的Codeblock 24以了解我如何實現該inference()函數。
# Codeblock 24
def inference():
denoised_images = [] #(1)
with torch.no_grad(): #(2)
current_prediction = torch.randn((64, NUM_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)).to(DEVICE) #(3)
for i in tqdm(reversed(range(NUM_TIMESTEPS))): #(4)
predicted_noise = model(current_prediction, torch.as_tensor(i).unsqueeze(0)) #(5)
current_prediction, denoised_image = noise_scheduler.backward_diffusion(current_prediction, predicted_noise, torch.as_tensor(i)) #(6)
if i%100 == 0: #(7)
denoised_images.append(denoised_image)
return denoised_images在標有#(1)的行中，我初始化了一個空列表，該列表將用于每100個時間步存儲一次去噪結果（#(7)）。這將使我們能夠稍后看到反向擴散的進行方式。實際的推理過程封裝在torch.no_grad()（#(2)）內。請記住，在擴散模型中，我們從完全隨機的噪聲中生成圖像，我們假設這些圖像最初是在t=999時。為了實現這一點，我們可以簡單地使用torch.randn()，如#(3)行所示。在這里，我們初始化一個大小為64×1×28×28的張量，表示我們即將同時生成64張圖像。接下來，我們編寫一個for從999開始向后迭代到0的循環（#(4)）。在這個循環中，我們將當前圖像和時間步作為訓練后的U-Net的輸入，并讓它預測噪聲（#(5)）。然后在#(6)行執行實際的反向擴散。在迭代結束時，我們應該得到類似于我們數據集中的新圖像。現在，讓我們在下面的代碼塊中調用該inference()函數。
# Codeblock 25
denoised_images = inference()推理完成后，我們現在可以看到生成的圖像是什么樣子。下面的Codeblock 26用于顯示我們剛剛生成的前42張圖像。
# Codeblock 26
fig, axes = plt.subplots(ncols=7, nrows=6, figsize=(10, 8))
counter = 0
for i in range(6):
for j in range(7):
axes[i,j].imshow(denoised_images[-1][counter].squeeze().detach().cpu().numpy(), cmap='gray') #(1)
axes[i,j].get_xaxis().set_visible(False)
axes[i,j].get_yaxis().set_visible(False)
counter += 1
plt.show()
圖18.在MNIST手寫數字數據集【引文3】上訓練的擴散模型生成的圖像
如果我們看一下上面的代碼塊，你會看到#(1)行的索引器[-1]表示我們只顯示來自最后一次迭代（對應于時間步長0）的圖像。這就是你在圖18中看到的圖像都沒有噪音的原因。我承認這可能不是最好的結果，因為并非所有生成的圖像都是有效的數字。——但是，這反而表明這些圖像不僅僅是原始數據集的重復。
這里我們還可以使用下面的Codeblock 27可視化反向擴散過程。你可以在圖19中的結果輸出中看到，我們最初從完全隨機的噪聲開始，隨著我們向右移動，噪聲逐漸消失。
# Codeblock 27
fig, axes = plt.subplots(ncols=10, figsize=(24, 8))
sample_no = 0
timestep_no = 0
for i in range(10):
axes[i].imshow(denoised_images[timestep_no][sample_no].squeeze().detach().cpu().numpy(), cmap='gray')
axes[i].get_xaxis().set_visible(False)
axes[i].get_yaxis().set_visible(False)
timestep_no += 1
plt.show()
圖19.時間步900、800、700等…直到時間步0時圖像的樣子【引文3】
結論
你可以基于本文提供的線路，進一步進行很多方面的研究。首先，如果你想要更好的結果，你可能需要調整Codeblock 2中的參數配置。其次，除了我們在下采樣和上采樣階段使用的卷積層堆棧之外，還可以通過實現注意層來改進U-Net模型。這樣做并不能保證你獲得更好的結果，尤其是對于像本文所使用的這樣的簡單數據集，但絕對值得一試。第三，如果你想挑戰自己，你還可以嘗試使用更復雜的數據集。
在實際應用中，擴散模型實際上可以做很多事情。最簡單的一個可能是數據增強。使用擴散模型，我們可以輕松地從特定的數據分布中生成新圖像。例如，假設我們正在進行一個圖像分類項目，但類別中的圖像數量不平衡。為了解決這個問題，我們可以從少數類別中取出圖像并將它們輸入到擴散模型中。通過這樣做，我們可以要求訓練后的擴散模型從該類別中生成任意數量的樣本。
好了，以上就是有關擴散模型的理論和實現的全部內容。感謝閱讀，希望你今天能學到一些新知識！
你可以通過此鏈接訪問該項目中使用的代碼。這里還有我之前關于自動編碼器、變分自動編碼器（VAE）、神經風格遷移（NST）和Transformer的文章的鏈接。
參考
【1】Jascha Sohl-Dickstein等人，《利用非平衡熱力學進行深度無監督學習》，Arxiv。
【2】Jonathan Ho等人，《去噪擴散概率模型》，Arxiv。
【3】Olaf Ronneberger等人，《U-Net：用于生物醫學圖像分割的卷積網絡》，Arxiv。
【4】Yann LeCun等人，《MNIST手寫數字數據庫》，Creative Commons Attribution-Share Alike 3.0許可證。
【5】Ashish Vaswani等人，《注意力就是你所需要的一切》，Arxiv。
譯者介紹
朱先忠，51CTO社區編輯，51CTO專家博客、講師，濰坊一所高校計算機教師，自由編程界老兵一枚。
原文標題：The Art of Noise，作者：Muhammad Ardi
