File size: 1,904 Bytes
c4b649d | 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 | """
model.py
========
Định nghĩa kiến trúc mô hình ResNet-50 cho 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 ResNet50_DR(nn.Module):
"""
Kiến trúc mô hình ResNet-50 cho phân loại 5 lớp DR:
- Backbone: ResNet-50 (Feature Extractor: 2048 channels)
- Classifier: Dropout(drop_rate) + Linear(2048 -> 5 classes)
"""
def __init__(
self,
num_classes: int = 5,
drop_rate: float = 0.3,
pretrained: bool = False,
):
super().__init__()
if pretrained:
weights = models.ResNet50_Weights.DEFAULT
self.model = models.resnet50(weights=weights)
else:
self.model = models.resnet50(weights=None)
in_features = self.model.fc.in_features # 2048 channels
self.model.fc = nn.Sequential(
nn.Dropout(p=drop_rate),
nn.Linear(in_features, num_classes),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.model(x)
def load_state_dict(self, state_dict: dict, strict: bool = True):
"""Hỗ trợ nạp cả state_dict của torchvision.resnet50 chuẩn lẫn wrapped module."""
if any(k.startswith("model.") for k in state_dict.keys()):
return super().load_state_dict(state_dict, strict=strict)
else:
return self.model.load_state_dict(state_dict, strict=strict)
def build_model(
num_classes: int = 5,
drop_rate: float = 0.3,
pretrained: bool = False,
) -> ResNet50_DR:
"""Hàm helper khởi tạo mô hình ResNet50_DR."""
model = ResNet50_DR(
num_classes=num_classes,
drop_rate=drop_rate,
pretrained=pretrained,
)
return model
|