Spaces:
Sleeping
Sleeping
Delete model.py
Browse files
model.py
DELETED
|
@@ -1,36 +0,0 @@
|
|
| 1 |
-
import torch
|
| 2 |
-
import torchvision
|
| 3 |
-
|
| 4 |
-
from torch import nn
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
def create_effnetb2_model(num_classes:int=102,
|
| 8 |
-
seed:int=42):
|
| 9 |
-
"""Creates an EfficientNetB2 feature extractor model and transforms.
|
| 10 |
-
|
| 11 |
-
Args:
|
| 12 |
-
num_classes (int, optional): number of classes in the classifier head.
|
| 13 |
-
Defaults to 3.
|
| 14 |
-
seed (int, optional): random seed value. Defaults to 42.
|
| 15 |
-
|
| 16 |
-
Returns:
|
| 17 |
-
model (torch.nn.Module): EffNetB2 feature extractor model.
|
| 18 |
-
transforms (torchvision.transforms): EffNetB2 image transforms.
|
| 19 |
-
"""
|
| 20 |
-
# Create EffNetB2 pretrained weights, transforms and model
|
| 21 |
-
weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
|
| 22 |
-
transforms = weights.transforms()
|
| 23 |
-
model = torchvision.models.efficientnet_b2(weights=weights)
|
| 24 |
-
|
| 25 |
-
# Freeze all layers in base model
|
| 26 |
-
for param in model.parameters():
|
| 27 |
-
param.requires_grad = False
|
| 28 |
-
|
| 29 |
-
# Change classifier head with random seed for reproducibility
|
| 30 |
-
torch.manual_seed(seed)
|
| 31 |
-
model.classifier = nn.Sequential(
|
| 32 |
-
nn.Dropout(p=0.3, inplace=True),
|
| 33 |
-
nn.Linear(in_features=1408, out_features=num_classes),
|
| 34 |
-
)
|
| 35 |
-
|
| 36 |
-
return model, transforms
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|