# 〈 Diffusion Model 論文研究與實作心得 Part.3 〉 模型訓練、照片修復與結果呈現 (Finale)
---
Tu 2023/3/27
## 一、前言
在前兩篇文章[〈 Diffusion Model 論文研究與實作心得 Part.1 〉 前言與圖片雜訊前處理](https://hackmd.io/@Tu32/B1-m6Tuai) 和 [〈 Diffusion Model 論文研究與實作心得 Part.2 〉 U-Net 模型架構介紹與實作](https://hackmd.io/hHXkc3OmTlCobpMS_1cGng?both)中，我完成了資料前處理與模型的搭建。因此在Part.3(最終篇)就要來進行模型的訓練和結果呈現。
## 二、模型訓練
我們可以參考一下ddpm作者的sudo code，這樣對實作的步驟有很大的幫助。
我們的模型輸出是預測圖片的雜訊(對，不是修復後的圖)，拿去和加在上面的雜訊進行比較。所以`get_loss`函數應該有三個參數，X_0，timestep和model。
```python
def get_loss(x_0, t, model):
pass
```
比較所需要的有三個東西
1. 某個timestep的x
2. 實際加上的雜訊
3. 模型預測的雜訊
1.和2.可以用[〈 Diffusion Model 論文研究與實作心得 Part.1 〉 前言與圖片雜訊前處理](https://hackmd.io/@Tu32/B1-m6Tuai) 裡定義的
```python
def get_index_from_list(vals, t, x_shape):
"""
Returns a specific index t of a passed list of values vals
while considering the batch dimension.
"""
batch_size = t.shape[0]
out = vals.gather(-1, t.cpu())
return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)
def forward_diffusion_sample(x_0, t, device="cpu"):
"""
Takes an image and a timestep as input and
returns the noisy version of it
"""
noise = torch.randn_like(x_0)
sqrt_alphas_cumprod_t = get_index_from_list(sqrt_alphas_cumprod, t, x_0.shape)
sqrt_one_minus_alphas_cumprod_t = get_index_from_list(
sqrt_one_minus_alphas_cumprod, t, x_0.shape
)
# mean + variance
return sqrt_alphas_cumprod_t.to(device) * x_0.to(device) \
+ sqrt_one_minus_alphas_cumprod_t.to(device) * noise.to(device), noise.to(device)
```
這個函數會回傳前兩點需要的東西。
```python
def get_loss(x_0, t, model):
x_t, noise = forward_diffuse_process(x_0, t)
```
而第三點需要的東西則要使用我們上次架構的U-net模型
```python
def get_loss(x_0, t, model):
x_t, noise = forward_diffuse_process(x_0, t)
noise_prediction = model(x_t, t)
```
最後對noise和noise_prediction進行比較就能得到Loss了，這邊選用L2 Loss
```python
def get_loss(x_0, t, model):
x_t, noise = forward_diffuse_process(x_0, t)
noise_prediction = model(x_t, t)
return F.l2_loss(noise, noise_prediction)
```
如此一來就能開始進行訓練了！
optimizer選用Adam，epochs先選用20 (colab的資源讓我一次只敢做這麼多QQ)
```python
from torch.optim import Adam
dataloader = DataLoader(data, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
optimizer = Adam(model.parameters(), lr=0.001)
epochs = 20 # Try more!
for epoch in range(epochs):
for step, batch in enumerate(dataloader):
optimizer.zero_grad()
t = torch.randint(0, T, (BATCH_SIZE,), device=device).long()
loss = get_loss(model, batch[0], t)
loss.backward()
optimizer.step()
if epoch % 5 == 0 and step == 0:
print(f"Epoch {epoch} | step {step:03d} Loss: {loss.item()} ")
```
訓練output：
```
Epoch 0 | step 000 Loss: 0.8118380904197693
Epoch 5 | step 000 Loss: 0.2767971158027649
Epoch 10 | step 000 Loss: 0.29156017303466797
Epoch 15 | step 000 Loss: 0.24683958292007446
Epoch 20 | step 000 Loss: 0.22735241055488586
```
這個專案的心臟，Diffusion Model正式訓練完成 (感動
## 三、圖片修復與成果呈現
說到底，我們模型輸出的終究只是對雜訊的預測，因此還需要一點點的數學才能將這個雜訊預測用於修復原圖。
還記得第一篇提到的q(Xt|Xt-1)嗎？那是用於破壞照片的forward process，而現在的backward process(修復圖片)ddpm的論文作者使用p(Xt-1|Xt)代表。
這部分牽涉到很複雜的數學(我也不太懂)，所以我就放一部份的筆記和[完整數學算式的連結](https://lilianweng.github.io/posts/2021-07-11-diffusion-models/#nice)
總之經過一點魔法我們能透過最底下框起來的式子計算出前一步timestep的圖。
論文中作者好像"**憑經驗**"省略了一堆數學還得到更好的結果，所以實作的部分就依照上面的sudo code就行了。
此外，這邊會用到第一篇定義的變數，我放在下面方便理解。
```python
# Define beta schedule
T = 300
betas = linear_beta_schedule(timesteps=T)
# Pre-calculate different terms for closed form
alphas = 1. - betas
alphas_cumprod = torch.cumprod(alphas, axis=0)
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)
#新定義的
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
```
```python
@torch.no_grad() #記得寫這行，在sample的時候才不會逆向傳遞梯度
def sample_timestep(x, t):
"""
給一個被破壞的圖片x和timestep，回傳修復後的圖片
"""
#這邊基本都是按照sudo code的算式
betas_t = get_index_from_list(betas, t, x.shape)
sqrt_one_minus_alphas_cumprod_t = get_index_from_list(
sqrt_one_minus_alphas_cumprod, t, x.shape
)
sqrt_recip_alphas_t = get_index_from_list(sqrt_recip_alphas, t, x.shape)
# Call model (current image - noise prediction)
model_mean = sqrt_recip_alphas_t * (
x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t
)
posterior_variance_t = get_index_from_list(posterior_variance, t, x.shape)
if t == 0:
return model_mean
else:
noise = torch.randn_like(x)
return model_mean + torch.sqrt(posterior_variance_t) * noise
@torch.no_grad()
def sample_and_plot_image():
#首先，生成隨機雜訊
img_size = IMG_SIZE
img = torch.randn((1, 3, img_size, img_size), device=device)
#這部分是用plt來呈現成果
plt.figure(figsize=(15,15))
plt.axis('off')
num_images = 10
stepsize = int(T/num_images)
#從第T個timestep修復到第0個
for i in range(0,T)[::-1]:
t = torch.full((1,), i, device=device, dtype=torch.long)
img = sample_timestep(img, t)
if i % stepsize == 0:
plt.subplot(1, num_images, int(i/stepsize+1))
show_tensor_image(img.detach().cpu()) #第一篇的函式
plt.show()
```
來看看epoch=80時候的成果：
呃...雖然有點抽象，但多少能看出類似臉、眼睛、頭髮的色塊，如果將epochs調高一點應該能得到更好的成果。
## 四、結語(系列總結)
第一次寫這種系列文，從資料前處裡到訓練模型，雖然省略了很多細節，很多地方可能做得不夠好，但我對自己踏出的第一步感到挺滿意的。
我之後可能會再寫一篇補充 (？，講講怎麼改造這個模型，讓他能產出更高畫質的圖片或變成prompt-to-image模型，又或者我搞了一張顯卡把epochs跑完再來看看成果之類的。一樣，都是後話了。
### 相關資料
https://www.youtube.com/watch?v=a4Yfz2FxXiY
https://www.youtube.com/watch?v=HoKDTa5jHvg&t=1338s
https://huggingface.co/blog/annotated-diffusion
https://arxiv.org/pdf/2102.09672.pdf
https://arxiv.org/pdf/1503.03585.pdf
https://arxiv.org/pdf/2006.11239.pdf
https://theaisummer.com/latent-variable-models/#reparameterization-trick
https://theaisummer.com/diffusion-models/
https://brohrer.mcknote.com/zh-
https://lilianweng.github.io/posts/2021-07-11-diffusion-models/#nice
###### tags: `AI` `Deep Learning` `Diffusion Model`
Sign in with Wallet
Connect another wallet