# 〈 Diffusion Model 論文研究與實作心得 Part.3 〉模型訓練、照片修復與結果呈現 (Finale)

一、前言

在前兩篇文章〈 Diffusion Model 論文研究與實作心得 Part.1 〉前言與圖片雜訊前處理和〈 Diffusion Model 論文研究與實作心得 Part.2 〉 U-Net 模型架構介紹與實作中，我完成了資料前處理與模型的搭建。因此在Part.3(最終篇)就要來進行模型的訓練和結果呈現。

二、模型訓練

我們可以參考一下ddpm作者的sudo code，這樣對實作的步驟有很大的幫助。

我們的模型輸出是預測圖片的雜訊(對，不是修復後的圖)，拿去和加在上面的雜訊進行比較。所以get_loss函數應該有三個參數，X_0，timestep和model。

def get_loss(x_0, t, model):
    pass

比較所需要的有三個東西

某個timestep的x
實際加上的雜訊
模型預測的雜訊

1.和2.可以用〈 Diffusion Model 論文研究與實作心得 Part.1 〉前言與圖片雜訊前處理裡定義的

def forward_diffuse_process(x_0, t):
    '''
    回傳第t個timestep的圖片和加上的雜訊
    '''
    noise = torch.randn_like(x_0) #回傳與X_0相同size的noise tensor，也就是reparameterization的epsilon
    sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[t]
    sqrt_oneminus_alphas_cumprod_t = sqrt_oneminus_alphas_cumprod[t]

    #element-wise的運算
    return sqrt_alphas_cumprod_t*x_0 + sqrt_oneminus_alphas_cumprod_t*noise, noise

這個函數會回傳前兩點需要的東西。

def get_loss(x_0, t, model):
    x_t, noise = forward_diffuse_process(x_0, t)

而3. 則需要使用我們上次架構的U-net模型

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

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)

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)代表。

這部分牽涉到很複雜的數學(我也不太懂)，所以我就放一部份的筆記和完整數學算式的連結

總之經過一點魔法我們能透過最底下框起來的式子計算出前一步timestep的圖。

論文中作者好像"憑經驗"省略了一堆數學還得到更好的結果，所以實作的部分就依照上面的sudo code就行了。

此外，這邊會用到第一篇定義的變數，我放在下面方便理解。

# 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)

@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時候的成果：