AI影像論文(03)：U-Net 論文筆記整理

近期在查看相關論文時時常看到U-Net這個演算法架構，但過去自己並沒有非常理解，因此想藉這個機會來逼自己把論文看完XD

🔹U-Net是什麼?

這篇論文是由弗賴堡大學的Olaf Ronneberger、Philipp Fischer 和 Thomas Brox 發表+，U-Net 的設計有兩個主要部分：一個是“收縮路徑”，用來理解圖像的大致情況；另一個是“擴展路徑”，用來精確地找出圖像中的細節位置。這樣的設計讓這個模型即使只有很少的圖像作為學習材料，也能夠很好地訓練，而且還用了很多技巧來提高它處理圖像的能力。

🔹U-Net架構

U-Net的架構有三大部分分別是收縮路徑（Contracting Path）、擴展路徑（Expansive Path）以及最終層（Final Layer），它的獨特之處在於其對稱結構以及能夠進行端到端的訓練，以及由於對稱性模型可以更好的利用上下文資訊，同時進行精準定位。

🐌收縮路徑（Contracting Path）

什麼是收縮路徑呢?簡單來說呢每次下採樣(池化)後，特徵的通道數量會加倍，這樣子有助於捕捉更複雜的特徵。這個過程有點類似傳統卷積神經網絡中的特徵提取部分。

而這一部分最主要是由多個捲積層和池化層（Max Pooling）組成，用於提取圖像中的特徵。

from tensorflow.keras.layers import Conv2D, MaxPooling2D, concatenate, Input

from tensorflow.keras.models import Model

def contracting_path(input_layer):

    # 第一個卷積塊

    conv1 = Conv2D(64, (3, 3), activation='relu', padding='same')(input_layer)

    conv1 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv1)

    pool1 = MaxPooling2D((2, 2))(conv1)

    # 第二個卷積塊

    conv2 = Conv2D(128, (3, 3), activation='relu', padding='same')(pool1)

    conv2 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv2)

    pool2 = MaxPooling2D((2, 2))(conv2)

    return pool2, conv1, conv2

🐌擴展路徑（Expansive Path）

那什麼又是擴展路徑呢?在擴展路徑中每一步上採樣的特徵圖會與收縮路徑中相對應的特徵圖進行拼接（Concatenation），這有助於恢復圖像中的細節和位置。

擴展路徑包含上採樣（Upsampling）操作和卷積操作，通過這些操作恢復圖像的空間尺寸。每次上採樣都會跟隨一個上捲積，並將特徵通道數量減半。

from tensorflow.keras.layers import Conv2DTranspose

def expansive_path(input_layer, concat_1, concat_2):

    # 上採樣加卷積（上卷積）    up1 = Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(input_layer)

    merge1 = concatenate([up1, concat_2], axis=3)

    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same')(merge1)

    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv3)

    # 第二次上採樣和合併    up2 = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(conv3)

    merge2 = concatenate([up2, concat_1], axis=3)

    conv4 = Conv2D(64, (3, 3), activation='relu', padding='same')(merge2)

    conv4 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv4)

    return conv4

🐌最終層（Final Layer）

最後，一個 1x1 的卷積層被用來將每個64組件的特徵向量映射到所需的類別數目，從而實現像素級的分類。

def final_layer(input_layer):

    output = Conv2D(1, (1, 1), activation='sigmoid')(input_layer)

    return output

🐌完整的U-Net模型

def unet_model(input_size=(256, 256, 1)):

    inputs = Input(input_size)

    pool2, conv1, conv2 = contracting_path(inputs)

    conv4 = expansive_path(pool2, conv1, conv2)

    outputs = final_layer(conv4)

    model = Model(inputs=inputs, outputs=outputs)

    return model# 創建模型model = unet_model()

model.summary()

🔹訓練過程

在訓練過程中，論文提到使用了Data Augmentation技術擴充訓練資料，可以讓model學習到”形變不變性(deformation invariance)”。

什麼是形變不變性?

這個概念指的是模型或算法能夠識別或處理在不同形變下的相同物體或特徵的能力。形變可以包括物體的旋轉、縮放、扭曲或其他幾何變化。

在實際應用中，形變不變性非常重要，因為它允許系統在不同的視角、尺寸或形狀變化時，依然能夠準確地識別物體。

而在論文中作者使用3*3網格的隨機位移向量來產生平滑型變，來增加強健性。

def random_deformation_grid(size, device):

    # 創建3x3的位移向量，值範圍在-0.2到0.2之間

    delta = torch.rand(1, 2, 3, 3, device=device) * 0.4 - 0.2

    # 生成均勻分布的座標網格

    grid = F.affine_grid(torch.eye(2, 3).unsqueeze(0) + delta, size, align_corners=False)

    return grid

def apply_deformation(image, grid):

    # 應用變形

    return F.grid_sample(image, grid, align_corners=False)

另外一個trick是使用帶有權重的loss，主要目的是為使同類別但有相互碰觸到的目標分割出來。

🔹U-net在Stable-diffusion之應用

近幾年最火紅的圖像生成模型絕對會是Stable-diffusion，而在Stable-diffusion有個蠻重要的部分就是U-net。

U-net在Stable-difussion中有什麼樣的應用呢?最主要有兩個：細節的捕捉與增強＆多尺度的特徵融合。

🐌細節的捕捉與增強

我們前面提到。

什麼是擴展路徑呢?在擴展路徑中每一步上採樣的特徵圖會與收縮路徑中相對應的特徵圖進行拼接（Concatenation），這有助於恢復圖像中的細節和位置。

我們可以使用U-net中的跳躍拼接來維持和增強圖像的細節。這些連結允許在直接使用編碼器中的特徵進而在解碼階段細化圖像。

🐌多尺度的特徵融合

U-Net 架構由編碼器和解碼器組成。在此設計中，編碼器將影像表示壓縮為較低的解析度。同時，解碼器將較低解析度的表示重建回原始的高解析度影像，旨在減少雜訊。

因此可以利用此特性去融合不同尺度的特徵，進而精確控制圖像局部特徵的細節。

參考資料：

論文：https://arxiv.org/pdf/1505.04597

Cook your First U-Net in PyTorch：https://towardsdatascience.com/cook-your-first-u-net-in-pytorch-b3297a844cf3

AI绘图Stable Diffusion中关键技术：U-Net的应用：https://cloud.tencent.com/developer/article/2397703