U-Net: Kiến trúc cơ bản cho bài toán phân đoạn ảnh

Chuỗi U-Net — Bài 1: nền tảng. U-Net là kiến trúc mạng nơ-ron tích chập được thiết kế cho bài toán phân đoạn ảnh (image segmentation), trong đó mô hình phải dự đoán nhãn cho từng pixel thay vì chỉ đưa ra một nhãn cho toàn bộ ảnh. Kiến trúc được Olaf Ronneberger, Philipp Fischer và Thomas Brox giới thiệu tại MICCAI 2015, ban đầu nhằm giải quyết các bài toán phân đoạn ảnh y sinh với lượng dữ liệu gán nhãn tương đối hạn chế [1].

:contentReference[oaicite:0]{index=0}

Điểm cốt lõi của U-Net không nằm ở một phép toán riêng lẻ mà ở cách kết hợp encoder để học ngữ cảnh, decoder để phục hồi độ phân giải và skip connection để đưa thông tin không gian chi tiết từ encoder sang decoder.

Tóm tắt nội dung

• U-Net biến bài toán phân đoạn ảnh thành phép ánh xạ từ ảnh đầu vào sang bản đồ lớp ở mức pixel.
• Encoder giảm dần kích thước không gian để học đặc trưng ngữ nghĩa; decoder tăng dần kích thước để khôi phục vị trí.
• Skip connection nối các mức tương ứng của encoder và decoder, giúp giữ biên và chi tiết nhỏ.
• U-Net không bắt buộc phải dùng Dice loss; phiên bản gốc sử dụng softmax theo pixel kết hợp cross-entropy có trọng số.

1. U-Net giải quyết bài toán gì?

Trong phân loại ảnh, mô hình có thể nhận một ảnh kích thước H × W × C và trả về một nhãn như “mèo”, “chó” hoặc “ô tô”. Phân đoạn ngữ nghĩa (semantic segmentation) khó hơn ở chỗ mỗi vị trí trong ảnh phải được gán một lớp. Vì vậy, đầu ra không còn là một vector nhỏ mà là một bản đồ dự đoán có cấu trúc không gian.

Bài toán	Đầu vào	Đầu ra	Mục tiêu
Phân loại ảnh	H × W × C	Một nhãn	Ảnh thuộc lớp nào?
Phân đoạn ngữ nghĩa	H × W × C	H × W × K	Mỗi pixel thuộc lớp nào?

Với K lớp cần phân đoạn, mô hình có thể tạo K giá trị dự đoán tại mỗi pixel. Ví dụ, với ảnh y tế gồm hai lớp “mô cần quan tâm” và “nền”, đầu ra cuối cùng có thể được chuyển thành một mặt nạ nhị phân. Với bài toán nhiều lớp như đường, ô tô, người đi bộ và nền, mỗi pixel được chọn lớp có xác suất cao nhất.

Trước U-Net, Fully Convolutional Network (FCN) đã chỉ ra rằng mạng tích chập có thể được huấn luyện end-to-end cho dự đoán dày đặc ở mức pixel và kết hợp đặc trưng sâu với đặc trưng có độ phân giải cao [2]. U-Net phát triển ý tưởng này theo một cấu trúc encoder-decoder gần đối xứng, đặc biệt nhấn mạnh việc truyền đặc trưng độ phân giải cao sang nhánh giải mã. :contentReference[oaicite:1]{index=1}

2. Kiến trúc U-Net: encoder, bottleneck và decoder

Hình dạng chữ U xuất hiện khi biểu diễn mạng theo độ phân giải: nhánh trái liên tục giảm kích thước đặc trưng, phần đáy chứa biểu diễn có độ phân giải thấp nhất, còn nhánh phải tăng kích thước trở lại. Hai nhánh được nối với nhau tại các mức tương ứng.

Hình 1. Sơ đồ khái quát U-Net được dựng lại từ kiến trúc và Figure 1 của công trình gốc [1]. Hình nguồn và các thông số kiến trúc đã được đối chiếu với trang dự án U-Net chính thức của University of Freiburg. :contentReference[oaicite:2]{index=2}

Encoder: nén không gian, tăng mức trừu tượng

Ở nhánh encoder của U-Net gốc, mỗi mức gồm hai phép tích chập 3×3, mỗi phép tích chập được theo sau bởi ReLU. Sau đó, max pooling 2×2 với stride 2 làm giảm một nửa chiều cao và chiều rộng. Mỗi lần downsampling, số kênh đặc trưng được tăng gấp đôi [1].

:contentReference[oaicite:3]{index=3}

Nếu bỏ qua ảnh hưởng của padding, quy luật trực quan có thể viết:

$$H_l\approx \frac{H_0}{2^l},W_l\approx \frac{W_0}{2^l},C_l\approx 2^l{C}_0$$

Trong đó l là mức downsampling. Ý nghĩa quan trọng hơn con số cụ thể là: độ phân giải giảm để mạng nhìn được vùng ngữ cảnh lớn hơn, trong khi số kênh tăng để biểu diễn nhiều kiểu đặc trưng hơn.

Bottleneck: đặc trưng cô đọng nhất

Phần thấp nhất của chữ U chứa feature map có độ phân giải không gian nhỏ nhưng số kênh lớn. Đây là nơi mô hình tổng hợp các đặc trưng đã trải qua nhiều tầng tích chập và pooling. Có thể hiểu bottleneck như biểu diễn ngữ nghĩa sâu nhất trước khi quá trình phục hồi độ phân giải bắt đầu.

Decoder: từ đặc trưng sâu trở lại bản đồ pixel

Nhánh decoder tăng dần kích thước feature map. Trong U-Net gốc, mỗi bước sử dụng upsampling kết hợp tích chập 2×2, giảm số kênh, nối với feature map tương ứng từ encoder rồi tiếp tục hai phép tích chập 3×3. Cuối mạng, tích chập 1×1 ánh xạ vector đặc trưng tại mỗi vị trí sang số lớp cần dự đoán [1].

:contentReference[oaicite:4]{index=4}

Chi tiết thường bị bỏ qua: U-Net năm 2015 dùng convolution không padding. Vì vậy feature map ở encoder phải được crop trước khi nối với decoder. Trong Figure 1 của paper, ảnh đầu vào 572×572 tạo bản đồ phân đoạn 388×388. Nhiều cài đặt hiện đại dùng padding để đầu vào và đầu ra có cùng kích thước; đó là một thay đổi thực dụng, không phải cấu hình chính xác của U-Net gốc.

3. Skip connection: vì sao U-Net giữ được chi tiết?

Pooling giúp mạng học ngữ cảnh nhưng đồng thời làm mất độ chính xác về vị trí. Một feature map sâu có thể trả lời tốt câu hỏi “vùng này có vẻ chứa một tế bào hay không”, nhưng thông tin về biên chính xác của tế bào có thể đã bị giảm sau nhiều lần downsampling.

U-Net giải quyết vấn đề đó bằng cách chuyển trực tiếp feature map từ encoder sang decoder ở cùng mức độ phân giải. Trong kiến trúc gốc, hai feature map được concatenate theo chiều kênh, không phải cộng phần tử như residual connection.

Luồng thông tin tại một mức decoder:

Feature sâu → Upsampling + Feature encoder → Concatenate → Convolution

Decoder nhờ đó nhận đồng thời hai loại thông tin: ngữ cảnh ngữ nghĩa từ các tầng sâu và chi tiết không gian từ tầng encoder tương ứng. Đây là lý do U-Net đặc biệt phù hợp với những trường hợp phải xác định chính xác đường biên, vật thể nhỏ hoặc cấu trúc mảnh.

Skip connection cũng giải thích vì sao chỉ “thu nhỏ ảnh rồi phóng lớn trở lại” chưa tạo thành U-Net. Nếu loại bỏ các kết nối ngang, decoder phải tái dựng chi tiết chỉ từ biểu diễn đã bị nén mạnh ở bottleneck.

4. Softmax, hàm mất mát và cách đánh giá

Xác suất theo từng pixel

Giả sử tại pixel x, tầng cuối tạo ra K logit. Với phân đoạn nhiều lớp, softmax chuyển các logit thành phân phối xác suất:

$$p_k\left(\mathbf{x}\right)=\frac{\exp \left(a_k\right(\mathbf{x}\left)\right)}{\sum _{j=1}^K\exp \left(a_j\right(\mathbf{x}\left)\right)}$$

Ở đây, a_k(x) là logit của lớp k tại pixel x và p_k(x) là xác suất sau softmax. Dự đoán cuối cùng thường lấy lớp có xác suất lớn nhất tại mỗi pixel.

Weighted cross-entropy trong U-Net gốc

Một nhầm lẫn phổ biến là xem Dice loss như hàm mất mát bắt buộc của U-Net. Trong công trình năm 2015, tác giả huấn luyện bằng softmax theo pixel kết hợp cross-entropy có trọng số. Dưới dạng hàm mất mát cần tối thiểu hóa, có thể viết:

:contentReference[oaicite:5]{index=5}

$$\mathcal{L}=-\sum _{\mathbf{x}\in \Omega }w\left(\mathbf{x}\right)\log \left(p_{y\left(\mathbf{x}\right)}\right(\mathbf{x}\left)\right)$$

Trọng số w(x) cho phép một số pixel đóng góp mạnh hơn vào loss. Với bài toán tách các tế bào chạm nhau, tác giả thiết kế weight map để nhấn mạnh vùng biên hẹp giữa hai đối tượng:

$$w\left(\mathbf{x}\right)=w_c\left(\mathbf{x}\right)+w_0·\exp (-\frac{{\left(d_1\right(\mathbf{x})+d_2(\mathbf{x}\left)\right)}^2}{2{\sigma }^2})$$

d₁ và d₂ là khoảng cách đến biên của hai tế bào gần nhất, còn w_c hỗ trợ cân bằng tần suất lớp. Thành phần exponential tăng trọng số tại vùng nằm gần hai biên, khiến lỗi ở vùng phân cách các đối tượng bị phạt mạnh hơn [1].

Dice dùng ở đâu?

Dice coefficient là thước đo rất phổ biến trong phân đoạn ảnh, đặc biệt khi foreground nhỏ hơn nhiều so với background. Với hai mặt nạ nhị phân A và B:

$$\mathrm{Dice}(A,B)=\frac{2|A\cap B|}{\left|A\right|+\left|B\right|}$$

Dice bằng 1 khi hai mặt nạ trùng nhau hoàn toàn và tiến về 0 khi mức chồng lấp giảm. Trong các hệ thống hiện đại, Dice có thể được dùng làm metric hoặc chuyển thành một loss khả vi. Tuy nhiên, đó là lựa chọn huấn luyện bổ sung chứ không phải thành phần định nghĩa kiến trúc U-Net.

5. Cài đặt U-Net tối giản bằng PyTorch

Đoạn mã dưới đây giữ ba ý tưởng chính của U-Net: encoder, decoder và skip connection. Để dễ chạy trên ảnh kích thước tùy ý, convolution dùng padding=1; vì vậy đây là phiên bản thực hành hiện đại, không phải bản sao chính xác kiến trúc unpadded năm 2015.

unet_minimal.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module): def **init**(self, in_ch, out_ch): super().**init**()
self.block = nn.Sequential(
nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
)

```
def forward(self, x):
    return self.block(x)
```

class UpBlock(nn.Module): def **init**(self, in_ch, skip_ch, out_ch): super().**init**()
self.up = nn.ConvTranspose2d(
in_ch, out_ch, kernel_size=2, stride=2
)
self.conv = DoubleConv(out_ch + skip_ch, out_ch)

```
def forward(self, x, skip):
    x = self.up(x)

    # Bảo đảm kích thước khớp khi ảnh không chia hết hoàn toàn.
    if x.shape[-2:] != skip.shape[-2:]:
        x = F.interpolate(
            x,
            size=skip.shape[-2:],
            mode="bilinear",
            align_corners=False,
        )

    x = torch.cat([skip, x], dim=1)
    return self.conv(x)
```

class UNet(nn.Module): def **init**(
self, in_channels=3, num_classes=2, base=32
): super().**init**()

```
    self.enc1 = DoubleConv(in_channels, base)
    self.enc2 = DoubleConv(base, base * 2)
    self.bottleneck = DoubleConv(base * 2, base * 4)

    self.pool = nn.MaxPool2d(2)

    self.up2 = UpBlock(base * 4, base * 2, base * 2)
    self.up1 = UpBlock(base * 2, base, base)

    self.head = nn.Conv2d(base, num_classes, kernel_size=1)

def forward(self, x):
    s1 = self.enc1(x)
    s2 = self.enc2(self.pool(s1))

    x = self.bottleneck(self.pool(s2))

    x = self.up2(x, s2)
    x = self.up1(x, s1)

    return self.head(x)
```

model = UNet(in_channels=3, num_classes=2)

x = torch.randn(1, 3, 256, 256)
logits = model(x)

print(logits.shape) # torch.Size([1, 2, 256, 256])

Tensor đầu vào có dạng [batch, channel, height, width]. Với một ảnh RGB 256×256 và hai lớp, đầu ra là [1, 2, 256, 256]: mỗi pixel có hai logit tương ứng hai lớp.

Không cần đặt softmax bên trong model khi huấn luyện phân đoạn nhiều lớp bằng torch.nn.CrossEntropyLoss, vì hàm này nhận trực tiếp logits. Đối với phân đoạn nhị phân dùng một output channel, cấu hình thường gặp là num_classes=1 cùng BCEWithLogitsLoss.

Điểm cần phân biệt: đoạn mã trên chỉ có hai mức downsampling để giữ ví dụ ngắn. U-Net gốc sâu hơn: kiến trúc trong paper có bốn bước max pooling, số kênh tăng từ 64 lên 1024 ở phần sâu nhất và tổng cộng 23 convolutional layers [1]. :contentReference[oaicite:6]{index=6}

6. Khi nào U-Net phù hợp và giới hạn cần biết?

U-Net là baseline hợp lý khi đầu ra cần giữ cấu trúc không gian gần với ảnh đầu vào. Ngoài ảnh y sinh — lĩnh vực mà kiến trúc được đề xuất và tiếp tục được phát triển cho nhiều tác vụ định lượng tế bào [3] — nguyên lý encoder-decoder có skip connection cũng được áp dụng rộng rãi cho nhiều bài toán phân đoạn khác. :contentReference[oaicite:7]{index=7}

U-Net đặc biệt đáng cân nhắc khi:

• cần dự đoán mask ở mức pixel;
• đường biên hoặc cấu trúc nhỏ có ý nghĩa;
• muốn một baseline CNN rõ ràng, dễ cài đặt và dễ kiểm soát;
• dữ liệu không đủ lớn để thiết kế một hệ thống quá phức tạp từ đầu;
• cần mở rộng sang nhiều biến thể encoder-decoder sau này.

Các giới hạn quan trọng:

• Tốn bộ nhớ: các feature map độ phân giải cao của encoder phải được giữ lại để dùng cho skip connection.
• Mất cân bằng lớp: nếu vùng foreground rất nhỏ, accuracy theo pixel có thể cao dù mask dự đoán kém; loss và metric phải được chọn phù hợp.
• Ngữ cảnh rất xa: CNN thuần túy tích lũy receptive field qua nhiều tầng, nhưng một số bài toán cần mô hình hóa quan hệ toàn cục mạnh hơn.
• Không có một U-Net duy nhất: padding, normalization, activation, số tầng, backbone, cách upsampling và loss có thể thay đổi đáng kể giữa các cài đặt.

Vì vậy, khi đọc kết quả thực nghiệm, câu hỏi đúng không chỉ là “có dùng U-Net hay không” mà còn phải kiểm tra phiên bản U-Net nào, input resolution bao nhiêu, loss nào được sử dụng, cách chia dữ liệu ra sao và metric được tính như thế nào.

7. Kết luận

U-Net giải quyết một mâu thuẫn trung tâm của phân đoạn ảnh: mạng cần giảm độ phân giải để học ngữ cảnh nhưng đồng thời phải giữ thông tin vị trí để tạo mask chính xác. Encoder đảm nhiệm phần ngữ nghĩa, decoder phục hồi không gian, còn skip connection kết hợp hai loại thông tin này tại từng mức.

Nắm chắc ba thành phần trên quan trọng hơn việc ghi nhớ từng số kênh của kiến trúc năm 2015. Từ nền tảng này có thể tiếp tục phân tích sâu hơn về convolution, receptive field, cơ chế upsampling, Dice loss, dữ liệu huấn luyện và các biến thể như 3D U-Net, Attention U-Net hay U-Net++.

Tài liệu tham khảo

[1] O. Ronneberger, P. Fischer, and T. Brox, “U-Net: Convolutional Networks for Biomedical Image Segmentation,” in Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015, Lecture Notes in Computer Science, vol. 9351, Springer, 2015, pp. 234–241, doi: 10.1007/978-3-319-24574-4_28. :contentReference[oaicite:8]{index=8}

[2] J. Long, E. Shelhamer, and T. Darrell, “Fully Convolutional Networks for Semantic Segmentation,” in Proc. IEEE Conf. Computer Vision and Pattern Recognition (CVPR), Boston, MA, USA, 2015, pp. 3431–3440, doi: 10.1109/CVPR.2015.7298965. :contentReference[oaicite:9]{index=9}

[3] T. Falk et al., “U-Net: deep learning for cell counting, detection, and morphometry,” Nature Methods, vol. 16, pp. 67–70, 2019, doi: 10.1038/s41592-018-0261-2. :contentReference[oaicite:10]{index=10}