arturomacias commited on
Commit
bcdc1d8
·
1 Parent(s): c7ee801

Upload model.py

Browse files
Files changed (1) hide show
  1. model.py +22 -0
model.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torchvision
3
+
4
+ from torch import nn
5
+
6
+ def create_effnetb2_model(num_classes: int=101,
7
+ seed: int=42):
8
+
9
+ weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
10
+ transforms = weights.transforms()
11
+ model = torchvision.models.efficientnet_b2(weights=weights)
12
+
13
+ for param in model.parameters():
14
+ param.requires_grad = False
15
+
16
+ torch.manual_seed(seed)
17
+ model.classifier = nn.Sequential(
18
+ nn.Dropout(p=0.3, inplace=True),
19
+ nn.Linear(in_features=1408, out_features=num_classes, bias=True)
20
+ )
21
+
22
+ return model, transforms