| """ |
| 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 |
| 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 |
|
|
|
|
