我們一起快速學會一個算法-UNet
作者：程序員小寒
UNet 模型由兩部分組成：編碼器（Contracting Path）和解碼器（Expanding Path），中間通過跳躍連接（Skip Connections）相連。
大家好，我是小寒。
今天給大家分享一個超強的算法模型，UNet
UNet 是一種專門用于圖像分割任務的卷積神經網絡（CNN）架構，最早由 Olaf Ronneberger 等人在 2015 年提出。
UNet 的名字來源于其結構的對稱性，類似于字母“U”。UNet 模型由于其優越的分割性能，被廣泛應用于各種圖像分割任務，如醫學圖像分割等。
圖片
Unet 模型架構
UNet 模型由兩部分組成：編碼器（Contracting Path）和解碼器（Expanding Path），中間通過跳躍連接（Skip Connections）相連。
編碼器（收縮路徑）
編碼器部分主要用于提取輸入圖像的特征。
它由一系列的卷積層、ReLU激活函數、最大池化層（Max Pooling）組成。
- 每個卷積層通常包含兩次卷積操作（使用 3x3 卷積核），每次卷積操作后接一個 ReLU 激活函數。
- 之后，采用一個 2x2 的最大池化層（Max Pooling）進行下采樣，以減少特征圖的空間維度。
- 每次下采樣后，特征圖的空間尺寸減小，而通道數增加，以提取更高層次的特征。
解碼器（擴展路徑）
解碼器部分用于恢復圖像的空間信息，最終輸出與輸入圖像相同大小的分割結果。
它由上采樣（up-sampling）操作和卷積層組成。
- 上采樣（Upsampling），通常通過反卷積將特征圖的空間分辨率逐步恢復。
- 上采樣后，通過跳躍連接（Skip Connection）將對應層的編碼器特征與解碼器特征拼接在一起，這樣可以保留輸入圖像的細節。
- 拼接后的特征圖經過兩次卷積操作（同樣使用 3x3 卷積核）和 ReLU 激活函數進行處理。
- 最終，經過逐步上采樣和卷積，恢復到與輸入圖像相同的分辨率。
跳躍連接 (Skip Connections)
在UNet中，跳躍連接將編碼器中每一層的輸出與解碼器中相應層的輸入相連，確保模型在還原圖像分辨率時保留更多的細節信息。
這種連接允許網絡在進行上采樣時參考編碼器部分的特征，從而更好地復原高分辨率特征。
UNet模型的優點
- 高效處理小樣本數據集
UNet 最初設計用于生物醫學圖像分割，具有高效利用小樣本數據集的能力。 - 精細的分割結果
通過跳躍連接，UNet 能夠很好地保留高分辨率的細節，使得分割結果更為精確。 - 靈活性強
UNet 結構簡單且有效，容易擴展和調整，適應不同類型的分割任務。
案例分享
下面是一個使用 PyTorch 實現 UNet 模型的代碼示例。這個示例展示了一個簡化版的UNet模型，并應用于圖像分割任務。
import torch
import torch.nn as nn
import torch.nn.functional as F
class UNet(nn.Module):
def __init__(self, in_channels, out_channels):
super(UNet, self).__init__()
# 編碼器部分
self.encoder1 = self.double_conv(in_channels, 64)
self.encoder2 = self.double_conv(64, 128)
self.encoder3 = self.double_conv(128, 256)
self.encoder4 = self.double_conv(256, 512)
# 最底部的卷積
self.bottleneck = self.double_conv(512, 1024)
# 解碼器部分
self.upconv4 = self.upconv(1024, 512)
self.decoder4 = self.double_conv(1024, 512)
self.upconv3 = self.upconv(512, 256)
self.decoder3 = self.double_conv(512, 256)
self.upconv2 = self.upconv(256, 128)
self.decoder2 = self.double_conv(256, 128)
self.upconv1 = self.upconv(128, 64)
self.decoder1 = self.double_conv(128, 64)
# 最終的1x1卷積，用于生成分割圖
self.final_conv = nn.Conv2d(64, out_channels, kernel_size=1)
def double_conv(self, in_channels, out_channels):
"""兩次卷積操作"""
return nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def upconv(self, in_channels, out_channels):
"""上采樣操作"""
return nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
def forward(self, x):
# 編碼器部分
enc1 = self.encoder1(x)
enc2 = self.encoder2(F.max_pool2d(enc1, kernel_size=2))
enc3 = self.encoder3(F.max_pool2d(enc2, kernel_size=2))
enc4 = self.encoder4(F.max_pool2d(enc3, kernel_size=2))
# Bottleneck
bottleneck = self.bottleneck(F.max_pool2d(enc4, kernel_size=2))
# 解碼器部分
dec4 = self.upconv4(bottleneck)
dec4 = torch.cat((dec4, self.crop_tensor(enc4, dec4)), dim=1)
dec4 = self.decoder4(dec4)
dec3 = self.upconv3(dec4)
dec3 = torch.cat((dec3, self.crop_tensor(enc3, dec3)), dim=1)
dec3 = self.decoder3(dec3)
dec2 = self.upconv2(dec3)
dec2 = torch.cat((dec2, self.crop_tensor(enc2, dec2)), dim=1)
dec2 = self.decoder2(dec2)
dec1 = self.upconv1(dec2)
dec1 = torch.cat((dec1, self.crop_tensor(enc1, dec1)), dim=1)
dec1 = self.decoder1(dec1)
# 最后的1x1卷積生成輸出
return self.final_conv(dec1)
def crop_tensor(self, encoder_tensor, decoder_tensor):
"""裁剪編碼器張量，使其與解碼器張量大小匹配"""
_, _, H, W = decoder_tensor.size()
encoder_tensor = self.center_crop(encoder_tensor, H, W)
return encoder_tensor
def center_crop(self, tensor, target_height, target_width):
"""中心裁剪函數"""
_, _, h, w = tensor.size()
crop_y = (h - target_height) // 2
crop_x = (w - target_width) // 2
return tensor[:, :, crop_y:crop_y + target_height, crop_x:crop_x + target_width]
# 使用示例
model = UNet(in_channels=1, out_channels=1) # 輸入和輸出均為1通道（例如用于灰度圖像）
input_image = torch.randn(1, 1, 572, 572) # 隨機生成一個輸入圖像
output = model(input_image)
print(output.shape)責任編輯：武曉燕
來源：
程序員學長
