foodvision / model.py
allispaul's picture
initial commit
3898f71
raw
history blame
1.14 kB
from typing import Tuple
import torch
from torch import nn
import torchvision
def create_effnetb3_model(num_classes: int = 101,
seed: int = 4,
) -> Tuple[nn.Module, torchvision.transforms.Compose]:
"""Create an EfficientNetB2 feature extractor model and transforms.
Args:
num_classes: Number of classes to use for classification (default 3).
seed: Random seed for reproducibility (default 4).
Returns:
A tuple (model, transforms) of the model and its image transforms.
"""
weights = torchvision.models.EfficientNet_B3_Weights.DEFAULT
transforms = weights.transforms()
model = torchvision.models.efficientnet_b3(weights=weights)
# Freeze parameters below the head
for param in model.parameters():
param.requires_grad = False
# Replace the classifier head with one of appropriate size for the problem
torch.manual_seed(seed)
model.classifier = nn.Sequential(
nn.Dropout(p=0.3, inplace=True),
nn.Linear(in_features=1536, out_features=num_classes)
)
return model, transforms