eel4759_classification

Comparison of different ImageNet-based CNN models for classifying diabetic retinopathy images
Log | Files | Refs | README

models.py (1461B)

      1 import torch.nn as nn
      2 from torchvision.models import (
      3     resnet50, ResNet50_Weights,
      4     efficientnet_b0, EfficientNet_B0_Weights,
      5     vit_b_16, ViT_B_16_Weights,
      6 )
      7 
      8 
      9 def create_model(model_name: str, num_classes: int = 5, pretrained: bool = True) -> nn.Module:
     10     """Create a pretrained model with the classification head replaced for num_classes.
     11 
     12     Supported model_name values: 'resnet50', 'efficientnet_b0', 'vit_b_16'
     13     All backbone layers are trainable (full fine-tuning).
     14     """
     15     if model_name == 'resnet50':
     16         weights = ResNet50_Weights.DEFAULT if pretrained else None
     17         model = resnet50(weights=weights)
     18         model.fc = nn.Linear(model.fc.in_features, num_classes)  # 2048 -> num_classes
     19 
     20     elif model_name == 'efficientnet_b0':
     21         weights = EfficientNet_B0_Weights.DEFAULT if pretrained else None
     22         model = efficientnet_b0(weights=weights)
     23         # classifier is Sequential(Dropout(0.2), Linear(1280, 1000))
     24         model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
     25 
     26     elif model_name == 'vit_b_16':
     27         weights = ViT_B_16_Weights.DEFAULT if pretrained else None
     28         model = vit_b_16(weights=weights)
     29         # heads is Sequential(head=Linear(768, 1000))
     30         model.heads.head = nn.Linear(model.heads.head.in_features, num_classes)
     31 
     32     else:
     33         raise ValueError(f"Unknown model: {model_name!r}. Choose from: resnet50, efficientnet_b0, vit_b_16")
     34 
     35     return model