MiladAbdollahi commited on
Commit
25dbaad
1 Parent(s): 8df1082

Upload model.py

Browse files
Files changed (1) hide show
  1. model.py +36 -0
model.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torchvision
3
+
4
+ from torch import nn
5
+
6
+
7
+ def create_effnetb2_model(num_classes:int=3,
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