Spaces:
Build error
Build error
import torchvision | |
from torch import nn | |
def create_effnetb0(num_classes:int=4, seed:int=42): | |
# 1. Get the base mdoel with pretrained weights and send to target device | |
weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT | |
transforms = weights.transforms() | |
model = torchvision.models.efficientnet_b0(weights=weights)#.to(device) | |
# 2. Freeze the base model layers | |
for param in model.features.parameters(): | |
param.requires_grad = False | |
# 3. Change the classifier head | |
model.classifier = nn.Sequential( | |
nn.Dropout(p=0.2), | |
nn.Linear(in_features=1280, out_features=num_classes) | |
)#.to(device) | |
# 5. Give the model a name | |
model.name = "effnetb0" | |
print(f"[INFO] Created new {model.name} model.") | |
return model, transforms | |