File size: 4,021 Bytes
47542cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """
model.py
========
Định nghĩa kiến trúc mô hình EfficientNet-B4 + CBAM (Convolutional Block Attention Module)
phục vụ bài toán phân loại Mức độ Bệnh Võng mạc Tiểu đường (Diabetic Retinopathy - 5 lớp ICDR).
"""
from __future__ import annotations
import torch
import torch.nn as nn
from torchvision import models
class ChannelAttention(nn.Module):
"""
Channel Attention Sub-module của CBAM.
Tính toán trọng số chú ý cho từng kênh đặc trưng dựa trên AvgPool và MaxPool.
"""
def __init__(self, in_planes: int, ratio: int = 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: torch.Tensor) -> torch.Tensor:
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):
"""
Spatial Attention Sub-module của CBAM.
Tập trung vào các vùng không gian quan trọng (xuất huyết, vi phình mạch, xuất tiết).
"""
def __init__(self, kernel_size: int = 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: torch.Tensor) -> torch.Tensor:
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):
"""
Convolutional Block Attention Module (CBAM).
Kết hợp Channel Attention và Spatial Attention nối tiếp.
"""
def __init__(self, in_planes: int, ratio: int = 16, kernel_size: int = 7):
super().__init__()
self.ca = ChannelAttention(in_planes, ratio)
self.sa = SpatialAttention(kernel_size)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x * self.ca(x)
x = x * self.sa(x)
return x
class EfficientNetB4_CBAM(nn.Module):
"""
Kiến trúc mô hình chính:
- Backbone: EfficientNet-B4 (Feature Extractor: 1792 channels)
- Attention: CBAM Module đặt ngay sau feature extractor
- Global Pooling: AdaptiveAvgPool2d(1)
- Classifier: Dropout(0.3) + Linear(1792 -> 5 classes)
"""
def __init__(
self,
num_classes: int = 5,
drop_rate: float = 0.3,
cbam_ratio: int = 16,
pretrained: bool = False,
):
super().__init__()
if pretrained:
weights = models.EfficientNet_B4_Weights.DEFAULT
backbone = models.efficientnet_b4(weights=weights)
else:
backbone = models.efficientnet_b4(weights=None)
self.features = backbone.features
in_planes = 1792 # Output channel size of EfficientNet-B4 features
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: torch.Tensor) -> torch.Tensor:
x = self.features(x)
x = self.cbam(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
return self.classifier(x)
def build_model(
num_classes: int = 5,
drop_rate: float = 0.3,
cbam_ratio: int = 16,
pretrained: bool = False,
) -> EfficientNetB4_CBAM:
"""Hàm helper khởi tạo mô hình và in số lượng tham số."""
model = EfficientNetB4_CBAM(
num_classes=num_classes,
drop_rate=drop_rate,
cbam_ratio=cbam_ratio,
pretrained=pretrained,
)
return model
|