SegFormer là gì? Transformer cho Semantic Segmentation

SegFormer là kiến trúc học sâu dành cho phân đoạn ngữ nghĩa (semantic segmentation), kết hợp một Transformer encoder phân cấp với decoder MLP rất nhẹ. Thay vì sử dụng Transformer ở một độ phân giải duy nhất như Vision Transformer (ViT), SegFormer tạo đặc trưng ở nhiều tỉ lệ không gian, phù hợp hơn với bài toán cần dự đoán nhãn cho từng pixel.

Điểm đáng chú ý của SegFormer không chỉ nằm ở việc đưa self-attention vào Computer Vision. Kiến trúc được thiết kế xoay quanh bốn ý tưởng: Mix Transformer (MiT), Overlapping Patch Embedding, Efficient Self-Attention, Mix-FFN và All-MLP Decoder. SegFormer cũng không sử dụng positional embedding cố định như ViT, nhờ đó tránh việc phải nội suy positional embedding khi độ phân giải ảnh thay đổi.

Tóm tắt nội dung

• SegFormer gồm MiT encoder phân cấp và All-MLP decoder.
• Encoder tạo bốn feature map ở các độ phân giải 1/4, 1/8, 1/16 và 1/32 so với ảnh đầu vào.
• Efficient Self-Attention giảm số lượng token dùng cho Key và Value, giúp attention trên ảnh có chi phí thấp hơn.
• Mix-FFN đưa depthwise convolution 3×3 vào FFN để bổ sung thông tin không gian mà không cần positional embedding cố định.

SegFormer trích xuất đặc trưng ở bốn mức phân giải bằng MiT encoder, sau đó đưa chúng về cùng tỉ lệ và hợp nhất bằng một decoder MLP nhẹ.

SegFormer giải quyết bài toán gì?

Trong phân loại ảnh, một ảnh thường được ánh xạ thành một nhãn duy nhất. Ngược lại, semantic segmentation yêu cầu mô hình dự đoán một lớp cho từng pixel.

Ảnh H × W × 3   →   SegFormer   →   Mask H × W

Ví dụ trong ảnh đường phố, mỗi pixel có thể được phân loại thành road, car, person, building hoặc sky. Trong ảnh y tế, các pixel có thể biểu diễn nền, cơ quan hoặc vùng tổn thương.

Điều này khác với instance segmentation. Semantic segmentation chỉ quan tâm pixel thuộc lớp nào. Nếu ảnh có ba chiếc ô tô, các pixel của cả ba vẫn mang cùng nhãn car. Instance segmentation còn phải xác định chúng là ba đối tượng riêng biệt.

Bài toán	Đầu ra	Ví dụ
Image Classification	Một hoặc nhiều nhãn cho toàn ảnh	Ảnh có mèo
Object Detection	Bounding box + class	Vị trí từng chiếc xe
Semantic Segmentation	Class cho từng pixel	Pixel road, car, sky...
Instance Segmentation	Mask cho từng instance	Xe 1, xe 2, xe 3

SegFormer được thiết kế cho semantic segmentation. Mô hình gốc được đánh giá trên các bộ dữ liệu như ADE20K, Cityscapes và COCO-Stuff.

Kiến trúc tổng thể: MiT Encoder và All-MLP Decoder

SegFormer có cấu trúc encoder-decoder, nhưng decoder đơn giản hơn đáng kể so với nhiều kiến trúc segmentation truyền thống. Hai thành phần chính là:

• Mix Transformer Encoder (MiT): tạo đặc trưng đa tỉ lệ từ ảnh.
• All-MLP Decoder: đưa các feature map về cùng không gian biểu diễn, hợp nhất chúng và dự đoán lớp cho từng vị trí.

Với ảnh đầu vào kích thước H × W × 3, MiT tạo ra bốn feature map:

F₁ ∈ R^{H/4 × W/4 × C₁}

F₂ ∈ R^{H/8 × W/8 × C₂}

F₃ ∈ R^{H/16 × W/16 × C₃}

F₄ ∈ R^{H/32 × W/32 × C₄}

Đây là một khác biệt quan trọng so với ViT nguyên bản. ViT thường tạo một chuỗi token ở một tỉ lệ duy nhất, trong khi semantic segmentation cần đồng thời giữ thông tin chi tiết về biên và thông tin ngữ nghĩa ở vùng rộng. MiT xây dựng một feature pyramid trực tiếp trong encoder.

Các stage đầu có độ phân giải cao nên giữ được nhiều thông tin không gian. Các stage sâu hơn có độ phân giải thấp nhưng biểu diễn ngữ nghĩa và ngữ cảnh rộng hơn. Decoder sau đó kết hợp cả bốn mức thay vì chỉ sử dụng feature map cuối cùng.

Cách hình dung: Stage đầu giúp biết chính xác biên đối tượng nằm ở đâu, trong khi stage sâu giúp hiểu vùng đó có khả năng là đối tượng gì. Semantic segmentation cần cả hai loại thông tin.

MiT Encoder: phần quan trọng nhất của SegFormer

Mix Transformer không đơn thuần là ViT được đưa vào segmentation. MiT sửa đổi cách tạo patch, self-attention và feed-forward network để phù hợp với ảnh độ phân giải cao.

Overlapping Patch Embedding

ViT nguyên bản chia ảnh thành các patch không chồng lấn. Cách này tiện cho việc biến ảnh thành chuỗi token nhưng có thể làm giảm tính liên tục tại ranh giới giữa các patch.

SegFormer sử dụng Overlapping Patch Embedding: vùng ảnh được lấy với kernel lớn hơn stride nên các patch lân cận có phần chồng lên nhau.

Stage	Kernel	Stride	Padding	Output
Stage 1	7 × 7	4	3	1/4
Stage 2–4	3 × 3	2	1	Giảm một nửa mỗi stage

Thiết kế này tạo lần lượt các feature map ở 1/4, 1/8, 1/16 và 1/32 kích thước ảnh gốc, đồng thời duy trì sự liên tục cục bộ tốt hơn patch không chồng lấn.

Efficient Self-Attention

Self-attention chuẩn là một vấn đề khi xử lý ảnh lớn. Với N token, attention có dạng:

Attention(Q,K,V) = Softmax(QK^T / √d) V

Ma trận QK^T có kích thước liên quan đến N × N, vì vậy chi phí của self-attention chuẩn tăng theo bậc hai của số token:

O(N²)

SegFormer giảm số phần tử không gian được sử dụng cho Key và Value trước khi thực hiện attention. Trực giác là Query vẫn cần đại diện cho các vị trí đang được cập nhật, nhưng Key và Value có thể được biểu diễn ở độ phân giải thấp hơn.

Chi tiết dễ gây nhầm: bài báo mô tả hệ số giảm độ dài chuỗi tương ứng với [64, 16, 4, 1] qua bốn stage. Trong nhiều implementation, tham số không gian được ghi là sr_ratio = [8, 4, 2, 1]; giảm mỗi chiều đi 8 lần tương ứng giảm số vị trí không gian khoảng 8² = 64 lần ở Stage 1. Cấu hình SegFormer hiện tại trong Transformers cũng biểu diễn theo sr_ratio này.

Nhờ giảm chuỗi K/V, mô hình vẫn khai thác được quan hệ xa trong ảnh nhưng tránh chi phí của full-resolution attention ở mọi stage.

Mix-FFN: đưa thông tin vị trí vào Transformer

Transformer cần biết vị trí của token. ViT giải quyết vấn đề này bằng positional embedding. SegFormer chọn một hướng khác: không dùng positional embedding cố định mà chèn depthwise convolution 3×3 vào feed-forward network.

x_out = MLP(GELU(DWConv_3×3(MLP(x_in)))) + x_in

Convolution 3×3 khai thác cấu trúc không gian cục bộ và padding cung cấp thông tin liên quan đến vị trí. Theo thiết kế của SegFormer, cơ chế này đủ để loại bỏ positional encoding cố định khỏi encoder.

Điều này đặc biệt hữu ích cho segmentation vì ảnh khi inference có thể có độ phân giải khác với ảnh lúc huấn luyện. Với positional embedding cố định, một số kiến trúc phải nội suy embedding sang kích thước mới; SegFormer tránh bước đó.

All-MLP Decoder hoạt động như thế nào?

Một điểm khác biệt lớn của SegFormer là decoder rất đơn giản. Thay vì xây dựng một decoder nhiều convolution, attention block hoặc pyramid module phức tạp, SegFormer trực tiếp hợp nhất bốn feature map của MiT bằng các phép chiếu tuyến tính.

Đầu tiên, mỗi feature map F_i được chiếu về cùng số chiều C:

F̂_i = Linear_Cᵢ→C(F_i)

Sau đó F̂₂, F̂₃ và F̂₄ được upsample về cùng độ phân giải với F̂₁, tức khoảng H/4 × W/4.

F = Concat(F̂₁, ↑F̂₂, ↑F̂₃, ↑F̂₄)

F_fused = Linear(F)

M = Linear_C→Nclasses(F_fused)

Decoder vì vậy kết hợp đặc trưng cục bộ của stage nông và thông tin ngữ cảnh của stage sâu mà không cần một decoder nặng. Trong mô hình gốc, đầu ra của decode head có độ phân giải H/4 × W/4 × N_classes.

Để thu được mask ở kích thước ảnh gốc, logits thường được nội suy lên H × W trước khi chọn lớp có giá trị lớn nhất tại từng pixel. Tài liệu Transformers hiện tại cũng lưu ý rằng logits của SegFormer không nhất thiết có cùng kích thước với ảnh đầu vào và cần resize trong bước hậu xử lý.

Từ logits đến segmentation mask

Giả sử logits sau khi resize có kích thước:

B × C × H × W

C là số lớp. Với mỗi pixel (h,w), mô hình chọn lớp có logit lớn nhất:

ŷ(h,w) = arg max_c z_c,h,w

Kết quả là một ma trận H × W, trong đó mỗi phần tử là ID của lớp dự đoán.

mIoU: chỉ số phổ biến để đánh giá segmentation

Một metric quan trọng của semantic segmentation là Intersection over Union (IoU). Với lớp c:

IoU_c = TP_c / (TP_c + FP_c + FN_c)

Mean IoU lấy trung bình trên các lớp được xét:

mIoU = (1/C) Σ_c=1^C IoU_c

Khác accuracy theo pixel, IoU phạt cả vùng dự đoán dư và vùng đối tượng bị bỏ sót, nên thường phản ánh chất lượng mask rõ hơn.

SegFormer khác U-Net, DeepLab và ViT ở đâu?

Không nên hiểu SegFormer đơn giản là “U-Net dùng Transformer”. Kiến trúc của hai mô hình có triết lý khác nhau.

Đặc điểm	U-Net	DeepLab	ViT cơ bản	SegFormer
Backbone	CNN	Thường CNN	Transformer	Hierarchical Transformer
Multi-scale encoder	Có	Có	Không phải đặc trưng của ViT gốc	Có, 4 stage
Ngữ cảnh xa	Tích lũy qua convolution	Atrous convolution / context modules	Self-attention toàn cục	Efficient self-attention
Decoder	Đối xứng, nhiều upsampling block	Context + decoder tùy biến thể	Không phải mô hình segmentation gốc	All-MLP nhẹ
Fixed positional embedding	Không cần	Không cần	Có trong ViT gốc	Không dùng

U-Net đặc biệt trực quan nhờ skip connection giữa encoder và decoder ở cùng tỉ lệ. DeepLab tập trung mạnh vào receptive field và atrous convolution. SegFormer thay phần trích xuất đặc trưng bằng Transformer phân cấp và đặt phần lớn sức mạnh biểu diễn vào encoder, nhờ đó decoder có thể giữ rất đơn giản.

Không có kiến trúc tốt nhất cho mọi dataset. Với dữ liệu nhỏ, ảnh y tế có cấu trúc đặc thù hoặc yêu cầu triển khai trên phần cứng hạn chế, CNN và hybrid CNN-Transformer vẫn có thể là lựa chọn tốt. SegFormer nên được đánh giá bằng validation set, tốc độ, bộ nhớ và yêu cầu triển khai thực tế.

Chạy SegFormer bằng Python

SegFormer đã được tích hợp trong thư viện Transformers. Ví dụ dưới đây sử dụng checkpoint SegFormer-B0 đã fine-tune trên ADE20K để thực hiện semantic segmentation cho một ảnh cục bộ.

import torch
from PIL import Image
from transformers import (
    AutoImageProcessor,
    SegformerForSemanticSegmentation,
)

model_name = "nvidia/segformer-b0-finetuned-ade-512-512"

processor = AutoImageProcessor.from_pretrained(model_name)
model = SegformerForSemanticSegmentation.from_pretrained(model_name)

image = Image.open("image.jpg").convert("RGB")

inputs = processor(
    images=image,
    return_tensors="pt"
)

model.eval()

with torch.no_grad():
    outputs = model(**inputs)

# Raw logits: B × C × H/4 × W/4
logits = outputs.logits

# Resize và chuyển logits thành class ID cho từng pixel
target_size = [(image.height, image.width)]

segmentation = processor.post_process_semantic_segmentation(
    outputs,
    target_sizes=target_size
)[0]

print("Logits:", logits.shape)
print("Mask:", segmentation.shape)
print(
    "Classes detected:",
    torch.unique(segmentation).tolist()
)

Đối với checkpoint 512 × 512 này, logits có dạng B × 150 × 128 × 128 sau preprocessing chuẩn, vì ADE20K sử dụng 150 lớp và decode head hoạt động ở khoảng 1/4 độ phân giải. Hàm hậu xử lý resize kết quả về kích thước ảnh yêu cầu trước khi trả về class ID cho từng pixel.

Nếu cần tên lớp thay vì ID, có thể đọc ánh xạ từ cấu hình của model:

class_ids = torch.unique(segmentation).tolist()

for class_id in class_ids:
    label = model.config.id2label[class_id]
    print(class_id, "→", label)

Khi fine-tune trên dataset riêng, số lớp của segmentation head phải được thay đổi theo số class của dữ liệu. Nếu bài toán chỉ có hai vùng như background và lesion, có thể dùng hai class logits hoặc thiết kế head/loss nhị phân phù hợp với pipeline riêng.

SegFormer-B0 đến B5 khác nhau thế nào?

SegFormer được xây dựng thành một họ mô hình từ B0 đến B5. Chúng giữ cùng nguyên lý kiến trúc nhưng tăng quy mô encoder. MiT-B0 là phiên bản nhẹ nhất, còn MiT-B5 lớn nhất trong họ mô hình gốc.

Biến thể	Quy mô tương đối	Phù hợp khi
B0	Nhỏ nhất	Ưu tiên tốc độ, bộ nhớ và thử nghiệm nhanh
B1–B2	Nhỏ đến trung bình	Cân bằng tài nguyên và năng lực mô hình
B3–B4	Lớn	GPU đủ mạnh và cần biểu diễn sâu hơn
B5	Lớn nhất trong họ gốc	Ưu tiên năng lực mô hình hơn chi phí tính toán

Không nên chọn B5 chỉ vì nó lớn hơn. Dataset nhỏ có thể không khai thác được lợi ích của mô hình lớn, trong khi batch size, thời gian huấn luyện và bộ nhớ GPU tăng đáng kể. Phiên bản phù hợp nên được chọn theo validation performance và ràng buộc triển khai.

Khi nào nên dùng SegFormer?

SegFormer đặc biệt phù hợp khi bài toán cần kết hợp thông tin cục bộ và ngữ cảnh xa. Điều này xuất hiện trong nhiều dạng dữ liệu ảnh như cảnh đường phố, ảnh vệ tinh, ảnh công nghiệp và ảnh y tế.

• Scene parsing: phân loại đường, trời, tòa nhà, cây, phương tiện và các thành phần của cảnh.
• Autonomous driving: xác định vùng có thể lái, phương tiện, người đi bộ và môi trường.
• Medical image segmentation: phân đoạn cơ quan, mô hoặc vùng bất thường khi có dữ liệu gán nhãn phù hợp.
• Remote sensing: phân loại đất, nước, công trình và lớp phủ bề mặt ở cấp pixel.
• Industrial vision: phân vùng bề mặt hoặc vùng lỗi trong hệ thống kiểm tra hình ảnh.

Tuy nhiên, SegFormer cũng có các giới hạn. Self-attention dù đã được giảm chi phí vẫn cần tài nguyên, đặc biệt khi ảnh rất lớn. Decoder nhẹ cũng không đồng nghĩa toàn bộ mô hình nhẹ, bởi phần lớn năng lực và chi phí nằm trong encoder.

Đối với segmentation có biên rất nhỏ hoặc đối tượng mảnh, chất lượng cuối cùng còn phụ thuộc vào độ phân giải đầu vào, augmentation, loss function và chiến lược upsampling. Không nên kỳ vọng chỉ đổi backbone sang Transformer sẽ tự động giải quyết các vấn đề mất cân bằng lớp hoặc mất chi tiết biên.

Một hiểu nhầm thường gặp: SegFormer không phải mô hình “không dùng convolution”. Overlapping Patch Embedding và đặc biệt Mix-FFN đều sử dụng convolution. Chính sự kết hợp giữa Transformer và các phép toán có tính cục bộ là một phần quan trọng của thiết kế.

Kết luận

SegFormer xây dựng semantic segmentation theo một hướng gọn nhưng hiệu quả: MiT encoder tạo đặc trưng phân cấp, Efficient Self-Attention giảm chi phí xử lý ảnh, Mix-FFN đưa thông tin không gian vào Transformer và All-MLP decoder hợp nhất các feature map đa tỉ lệ. Kiến trúc không phụ thuộc vào positional embedding cố định và đặt phần lớn năng lực biểu diễn vào encoder.

Điểm quan trọng nhất khi học SegFormer không phải ghi nhớ tên các block, mà hiểu luồng dữ liệu: ảnh → overlapping patch embedding → bốn stage MiT → feature pyramid → MLP projection → upsampling và fusion → per-pixel logits → segmentation mask. Khi nắm được luồng này, việc fine-tune SegFormer, thay decoder, kết hợp với CNN hoặc áp dụng knowledge distillation cho segmentation trở nên dễ phân tích hơn.