Cấu trúc Mô hình AI — EfficientNet-B4 + CBAM (Diabetic Retinopathy Classification)

Tài liệu chi tiết về kiến trúc mô hình AI, cấu hình huấn luyện, dòng chảy dữ liệu (data flow) và hướng dẫn đóng gói suy luận (inference packaging) cho bài toán Phân loại Mức độ Bệnh Võng mạc Tiểu đường (DR - 5 lớp ICDR).


1. Tổng quan Mô hình

Thông số Giá trị
Tên mô hình EfficientNetB4_CBAM
Bài toán Phân loại 5 mức độ bệnh Võng mạc Tiểu đường (ICDR Standard)
Backbone EfficientNet-B4 (Pretrained ImageNet)
Cơ chế Chú ý (Attention) CBAM (Convolutional Block Attention Module)
Số tham số (Parameters) ~17.5M (Backbone: 17.5M, CBAM: ~0.04M)
Kích thước đầu vào 3 × 224 × 224 (hoặc 3 × 380 × 380)
Đầu ra 5 logits / xác suất (Probabilities via Softmax)
Hàm mất mát (Loss Function) Focal Loss ($\gamma = 2.0$, Class-weighted)
Đánh giá chính Quadratic Weighted Kappa (QWK), Macro F1-Score, ROC-AUC

2. Các Lớp Đầu Ra (Target Classes)

Mô hình dự đoán 5 mức độ tổn thương võng mạc tiểu đường theo tiêu chuẩn quốc tế ICDR:

Class ID Tên nhãn (Label Name) Mô tả lâm sàng Tỷ lệ dữ liệu train
0 No DR Không phát hiện tổn thương DR 71.6%
1 Mild Xuất hiện vi phình mạch (Microaneurysms) 7.0%
2 Moderate Xuất hiện xuất huyết / xuất tiết nhẹ-vừa 15.8%
3 Severe Xuất huyết nặng 4 góc phần tư / chuỗi tĩnh mạch 2.8%
4 Proliferative DR Tăng sinh tân mạch / xuất huyết dịch kính 2.9%

3. Kiến trúc Chi tiết (Model Architecture)

Mô hình kết hợp giữa khả năng trích xuất đặc trưng đa tỷ lệ của EfficientNet-B4 và cơ chế lọc đặc trưng không gian & kênh của CBAM Attention.

3.1. Sơ đồ Dòng Chảy Dữ Liệu (Data Flow Diagram)

flowchart TD
    A["Input Image (3 x 224 x 224)"] --> B["EfficientNet-B4 Backbone"]
    B --> C["Feature Maps (1792 x 7 x 7)"]
    
    subgraph CBAM ["CBAM Attention Module"]
        C --> D["Channel Attention Module (CA)"]
        D -->|Feature * CA Map| E["Spatial Attention Module (SA)"]
        E -->|Feature * SA Map| F["Refined Features (1792 x 7 x 7)"]
    end
    
    F --> G["AdaptiveAvgPool2d (1 x 1)"]
    G --> H["Flatten (1792)"]
    H --> I["Dropout (p=0.3)"]
    I --> J["Linear Classification Head (1792 -> 5)"]
    J --> K["Logits Output (5)"]
    K --> L["Softmax (Probabilities)"]

3.2. Mô tả Chi tiết từng Thành phần

A. Backbone (EfficientNet-B4)

  • Chức năng: Trích xuất đặc trưng cấp cao từ ảnh đáy mắt (fundus).
  • Đầu ra trích xuất: Tensor có kích thước (Batch, 1792, H/32, W/32) — tương ứng với (Batch, 1792, 7, 7) khi đầu vào là 224x224.

B. CBAM Attention Module

Bao gồm 2 sub-modules nối tiếp nhau:

  1. Channel Attention (Chú ý theo Kênh):

    • Gom thông tin không gian bằng AdaptiveAvgPool2d(1)AdaptiveMaxPool2d(1).
    • Đưa qua shared MLP (2 lớp Conv2d giảm chiều theo tỷ lệ ratio=16, tức $1792 \rightarrow 112 \rightarrow 1792$).
    • Kết hợp kết quả bằng phép cộng, qua hàm Sigmoid tạo Trọng số Kênh $\mathbf{M}_c$.
    • Nhân phần tử với feature map ban đầu.
  2. Spatial Attention (Chú ý theo Không gian):

    • Gom thông tin kênh bằng phép mean(dim=1)max(dim=1) tạo tensor 2 kênh.
    • Đưa qua lớp Conv2d(2 -> 1, kernel_size=7, padding=3) và hàm Sigmoid tạo Trọng số Không gian $\mathbf{M}_s$.
    • Nhân phần tử để tập trung vào các vùng tổn thương quan trọng (xuất huyết, vi phình mạch, xuất tiết).

C. Classifier Head

  • AdaptiveAvgPool2d(1): Thu gọn đặc trưng không gian về 1792 x 1 x 1.
  • Flatten: Biến đổi thành vector 1D chiều dài 1792.
  • Dropout(p=0.3): Giảm thiểu hiện tượng Overfitting.
  • Linear(1792, 5): Ánh xạ đặc trưng về 5 điểm số đại diện cho 5 mức độ bệnh.

4. Cấu trúc File Checkpoint (.pth)

File checkpoint modelAI/efficientnet_b4_cbam_fold1.pth lưu trữ dưới dạng PyTorch dictionary:

checkpoint = {
    "epoch": 24,                                 # Epoch đạt kết quả tốt nhất
    "model_state_dict": model.state_dict(),     # Trọng số tất cả các layer
    "optimizer_state_dict": optimizer.state_dict(),
    "val_qwk": 0.8542,                           # QWK trên tập Validation
    "val_f1": 0.7215,                            # Macro F1 trên tập Validation
    "args": CONFIG                               # Dictionary lưu tham số cấu hình
}

5. Hướng dẫn Đóng gói Mô hình Suy luận (Inference Packaging Guide)

Để đưa mô hình này vào ứng dụng sản xuất (Production / Web / Mobile / REST API), bạn cần đóng gói bộ file theo cấu trúc chuẩn bên dưới.

5.1. Cấu trúc thư mục đóng gói đề xuất

modelAI/
├── README.md                              # Tài liệu cấu hình & kiến trúc (File này)
├── efficientnet_b4_cbam_fold1.pth         # File trọng số PyTorch checkpoint
├── phase_a_efficientnet_cbam_focal.ipynb  # Notebook huấn luyện & thực nghiệm
│
├── package/                               # Thư mục đóng gói suy luận (Inference SDK)
│   ├── config.json                        # Cấu hình nhãn & tiền xử lý
│   ├── model.py                           # Đã trích xuất PyTorch class (EfficientNetB4_CBAM)
│   ├── preprocessing.py                   # Đã trích xuất pipeline tiền xử lý ảnh
│   ├── predictor.py                       # Class DRPredictor chính để gọi suy luận
│   ├── requirements.txt                   # Danh sách thư viện phụ thuộc
│   └── test_inference.py                  # Script test suy luận thử 1 ảnh

5.2. Mã nguồn mã hóa Python Lớp Model (package/model.py)

import torch
import torch.nn as nn
from torchvision import models

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc = nn.Sequential(
            nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False),
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc(self.avg_pool(x))
        max_out = self.fc(self.max_pool(x))
        return self.sigmoid(avg_out + max_out)

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        return self.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1)))

class CBAM(nn.Module):
    def __init__(self, in_planes, ratio=16, kernel_size=7):
        super().__init__()
        self.ca = ChannelAttention(in_planes, ratio)
        self.sa = SpatialAttention(kernel_size)

    def forward(self, x):
        x = x * self.ca(x)
        x = x * self.sa(x)
        return x

class EfficientNetB4_CBAM(nn.Module):
    def __init__(self, num_classes=5, drop_rate=0.3, cbam_ratio=16):
        super().__init__()
        backbone = models.efficientnet_b4(weights=None)
        self.features = backbone.features
        in_planes = 1792
        self.cbam = CBAM(in_planes, ratio=cbam_ratio)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Dropout(p=drop_rate),
            nn.Linear(in_planes, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.cbam(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)

5.3. File Cấu hình Json (package/config.json)

{
  "model_name": "EfficientNetB4_CBAM",
  "num_classes": 5,
  "input_size": [224, 224],
  "mean": [0.485, 0.456, 0.406],
  "std": [0.229, 0.224, 0.225],
  "labels": {
    "0": "No DR",
    "1": "Mild",
    "2": "Moderate",
    "3": "Severe",
    "4": "Proliferative DR"
  }
}

5.4. Class Gọi Dự đoán (package/predictor.py)

import os
import json
import torch
from PIL import Image
from torchvision import transforms
from model import EfficientNetB4_CBAM

class DRPredictor:
    def __init__(self, weights_path="../efficientnet_b4_cbam_fold1.pth", config_path="config.json"):
        with open(config_path, "r", encoding="utf-8") as f:
            self.config = json.load(f)

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # 1. Khởi tạo kiến trúc
        self.model = EfficientNetB4_CBAM(
            num_classes=self.config["num_classes"],
            drop_rate=0.3
        )

        # 2. Load trọng số (hỗ trợ cả dạng checkpoint dict hoặc raw state_dict)
        checkpoint = torch.load(weights_path, map_location=self.device)
        if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
            self.model.load_state_dict(checkpoint["model_state_dict"])
        else:
            self.model.load_state_dict(checkpoint)

        self.model.to(self.device)
        self.model.eval()

        # 3. Pipeline Transform chuẩn化
        self.transform = transforms.Compose([
            transforms.Resize(tuple(self.config["input_size"])),
            transforms.ToTensor(),
            transforms.Normalize(mean=self.config["mean"], std=self.config["std"])
        ])

    def predict(self, image_path):
        """
        Nhận vào đường dẫn ảnh đáy mắt và trả về kết quả phân loại DR.
        """
        image = Image.open(image_path).convert("RGB")
        tensor_img = self.transform(image).unsqueeze(0).to(self.device)

        with torch.no_grad():
            outputs = self.model(tensor_img)
            probs = torch.softmax(outputs, dim=1)[0]
            pred_class = torch.argmax(probs).item()

        return {
            "class_id": pred_class,
            "class_name": self.config["labels"][str(pred_class)],
            "confidence": float(probs[pred_class]),
            "probabilities": {
                self.config["labels"][str(i)]: float(probs[i])
                for i in range(len(probs))
            }
        }

6. Yêu cầu Môi trường (Dependencies)

File package/requirements.txt:

torch>=2.0.0
torchvision>=0.15.0
Pillow>=9.5.0
numpy>=1.24.0
opencv-python>=4.7.0
Downloads last month
46
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Evaluation results