Jasur05 commited on
Commit
e9d2555
·
verified ·
1 Parent(s): 313d8dd

upload model.py

Browse files
Files changed (1) hide show
  1. model.py +24 -0
model.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torchvision
3
+ from torchvision.models import EfficientNet_B2_Weights
4
+ from torch import nn
5
+
6
+ def create_model(num_classes=7):
7
+ weights = EfficientNet_B2_Weights.DEFAULT
8
+ model = torchvision.models.efficientnet_b2(weights=weights)
9
+
10
+ for param in model.parameters():
11
+ param.requires_grad = False # Freeze for inference
12
+
13
+ model.classifier = nn.Sequential(
14
+ nn.Dropout(p=0.3),
15
+ nn.Linear(model.classifier[1].in_features, num_classes)
16
+ )
17
+
18
+ return model
19
+
20
+ def load_model(weights_path="fine_tuned_effnetb2_dermamnist.pth"):
21
+ model = create_model(num_classes=7)
22
+ model.load_state_dict(torch.load(weights_path, map_location=torch.device("cpu")))
23
+ model.eval()
24
+ return model