U-Net

U-Net là mạng nơ-ron tích chập (CNN) dành cho phân đoạn ngữ nghĩa, tác vụ gán một lớp cho từng pixel trong ảnh. Olaf Ronneberger, Philipp Fischer và Thomas Brox thuộc Đại học Freiburg giới thiệu U-Net năm 2015 trong bài báo U-Net: Convolutional Networks for Biomedical Image Segmentation. Encoder thu nhỏ ảnh để nắm bắt nội dung; decoder phóng kết quả về độ phân giải đầy đủ; các kết nối tắt truyền chi tiết ở từng tỷ lệ từ encoder sang decoder, giúp đường biên đối tượng sắc nét. U-Net được huấn luyện hiệu quả chỉ với vài chục ảnh hiển vi có nhãn, chiến thắng cách biệt ở hai hạng mục của thử thách theo dõi tế bào ISBI 2015 và vẫn là model tiêu chuẩn trong phân tích ảnh y tế. Sau này, thiết kế của U-Net trở thành mạng khử nhiễu trong nhiều model khuếch tán.

Cách U-Net hoạt động#

Mạng có hai nửa đối xứng, được minh họa như hai nhánh của chữ U:

• Nhánh thu hẹp (encoder): Các cặp phép tích chập 3×3 với activation ReLU được lặp lại; sau mỗi cặp là phép gộp cực đại 2×2 làm giảm một nửa độ phân giải. Số kênh tăng gấp đôi ở mỗi cấp, từ 64 lên 1.024, giúp mạng đánh đổi chi tiết không gian để có trường tiếp nhận rộng hơn và đặc trưng phong phú hơn.
• Nhánh mở rộng (decoder): Ở mỗi cấp, phép tích chập tăng mẫu 2×2 làm độ phân giải tăng gấp đôi và số kênh giảm một nửa; kết quả được nối với feature map của encoder ở cùng cấp, rồi hai phép tích chập 3×3 kết hợp chúng. Phép tích chập 1×1 cuối cùng ánh xạ đặc trưng của mỗi pixel thành điểm số lớp.

Mạng gốc có 23 lớp tích chập và bốn bước gộp. Mạng dùng phép tích chập không đệm, vì vậy tile đầu vào 572×572 tạo ra bản đồ phân đoạn 388×388; ảnh lớn được phân đoạn thành các tile chồng lấn. Hầu hết triển khai hiện nay đệm phép tích chập để đầu ra khớp với kích thước đầu vào.

Kết nối tắt#

Phép pooling làm mất thông tin về vị trí chính xác của từng đặc trưng. Tại nút thắt cổ chai, mạng biết một vùng chứa màng tế bào nhưng không biết màng đi qua những pixel nào. Các kết nối tắt sao chép bản đồ đặc trưng có độ phân giải cao của mỗi cấp encoder qua nhánh U sang decoder, nơi kết hợp thông tin "là gì" từ nhánh đã được upsample với thông tin "ở đâu" từ encoder. U-Net ghép nối hai phần theo chiều kênh, trong khi các kết nối dư của ResNet cộng các đặc trưng bên trong một khối. Đoạn mã PyTorch này xây dựng một cấp encoder, một nút thắt cổ chai và một cấp decoder:

import torch
from torch import nn

# One encoder level, a bottleneck, and one decoder level of a U-Net
encoder = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.ReLU())
bottleneck = nn.Sequential(nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())
upsample = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
decoder = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.ReLU())
head = nn.Conv2d(16, 2, kernel_size=1)  # 2 classes, e.g. cell and background

image = torch.randn(1, 3, 128, 128)
skip = encoder(image)  # full-resolution features kept for the decoder
context = bottleneck(skip)  # half resolution, twice the channels, wider context

# Skip connection: concatenate upsampled context with the saved high-resolution features
merged = torch.cat([upsample(context), skip], dim=1)
logits = head(decoder(merged))
print(f"context {tuple(context.shape)} -> merged {tuple(merged.shape)} -> logits {tuple(logits.shape)}")

Đoạn mã in context (1, 32, 64, 64) -> merged (1, 32, 128, 128) -> logits (1, 2, 128, 128): 16 kênh ngữ cảnh được tăng mẫu cùng 16 kênh kết nối tắt được xếp chồng thành 32 kênh; đầu ra chứa hai điểm số lớp cho mỗi pixel trong 128×128 pixel của đầu vào.

Huấn luyện với ít dữ liệu#

Tập huấn luyện ban đầu về cấu trúc neuron trong ảnh hiển vi điện tử chứa 30 ảnh kích thước 512×512 pixel. U-Net bù đắp lượng dữ liệu ít bằng tăng cường dữ liệu mạnh, đặc biệt là các biến dạng đàn hồi ngẫu nhiên mô phỏng cách mô uốn cong, cùng hàm mất mát entropy chéo có trọng số nhấn mạnh các khe nền mỏng giữa những tế bào tiếp xúc, nhờ đó các tế bào lân cận được phân tách thành những vùng riêng biệt. Model phân đoạn y tế hiện thường được huấn luyện bằng hàm mất mát Dice, đo trực tiếp mức độ chồng lấp giữa các mask và xử lý tốt cấu trúc nhỏ chỉ chiếm vài phần trăm số pixel.

Các biến thể và kiến trúc hậu duệ#

• 3D U-Net và V-Net: Thay các phép toán 2D bằng phép toán 3D để phân đoạn thể tích như ảnh chụp CT và MRI.
• Attention U-Net và UNet++: Bổ sung các cổng attention hoặc các đường kết nối tắt lồng nhau, kết nối dày đặc để tinh chỉnh những đặc trưng encoder được truyền đến decoder.
• nnU-Net: Framework tự động cấu hình tiền xử lý, kiến trúc và huấn luyện cho tập dữ liệu mới. Một nghiên cứu MICCAI 2024 nhận thấy các U-Net dựa trên CNN được tinh chỉnh tốt trong framework này vẫn vượt trội hơn các kiến trúc Transformer và Mamba mới hơn trong phân đoạn y tế 3D khi được đánh giá nghiêm ngặt.
• U-Net khuếch tán: DDPM và Stable Diffusion 1.x dùng U-Net để dự đoán nhiễu trong ảnh hoặc latent, đồng thời bổ sung các lớp attention để prompt điều hướng quá trình tạo sinh; U-Net của Stable Diffusion 1.x có khoảng 860 triệu tham số, còn ControlNet kết nối với encoder của U-Net này. Các bộ tạo sinh mới hơn thay U-Net bằng Transformer khuếch tán.

U-Net và các model liên quan#

• Mạng tích chập hoàn toàn (FCN): FCN năm 2014 chuyển bộ phân loại thành bộ dự đoán theo từng pixel và kết hợp dự đoán thô, chi tiết bằng cách cộng các bản đồ điểm số lớp. U-Net xây dựng decoder hoàn chỉnh, nối đặc trưng ở mọi cấp.
• Autoencoder: Cũng có cấu trúc encoder và decoder, nhưng được huấn luyện để tái tạo đầu vào qua một nút thắt hẹp. Kết nối tắt sẽ cho phép autoencoder bỏ qua nút thắt, trong khi U-Net cần chúng để tạo nhãn theo từng pixel được căn chỉnh.
• Model phân đoạn instance: U-Net xuất một lớp cho mỗi pixel, vì vậy hai đối tượng cùng lớp tiếp xúc với nhau sẽ chung nhãn, trừ khi quá trình huấn luyện hoặc hậu xử lý tách chúng ra.

Model phân đoạn ngữ nghĩa của Ultralytics YOLO26, có hậu tố -sem, sử dụng phiên bản rút gọn của decoder U-Net. Neck tăng mẫu đầu ra backbone ở tỷ lệ 1/32 hai lần, lần lượt nối đặc trưng backbone ở tỷ lệ 1/16 và 1/8; sau đó, model phân loại pixel ở tỷ lệ 1/8 và tăng mẫu điểm số lớp lên 8× để đạt kích thước đầu vào.

Ứng dụng thực tế#

• Phân tích ảnh vệ tinh: Bản đồ lớp phủ mặt đất gán nhãn từng pixel là nước, rừng, đất canh tác hoặc công trình để theo dõi nạn phá rừng và sự phát triển đô thị.
• Kiểm tra công nghiệp: Mask pixel của vết xước, vết nứt và ăn mòn đo chính xác kích thước và hình dạng khuyết tật, điều mà bounding box không thể làm được.
• Chẩn đoán hình ảnh y tế: U-Net phác thảo cơ quan, khối u và tổn thương trong ảnh chụp CT, MRI để hỗ trợ chẩn đoán và lập kế hoạch xạ trị, đồng thời tách từng tế bào trong ảnh hiển vi. Trong thử thách theo dõi tế bào ISBI 2015, U-Net được huấn luyện trên 20 ảnh tế bào HeLa có chú thích một phần đạt IoU trung bình 77,5%, so với 46% của phương pháp tốt thứ hai.