EngMohamed2023 commited on
Commit
3ef65f1
·
1 Parent(s): e5d81b6

Upload model.py

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