aayanb09 commited on
Commit
6d711a0
·
verified ·
1 Parent(s): 7c2f643

Create 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 torch.nn as nn
3
+ from torchvision import models
4
+
5
+ class FoodIngredientClassifier(nn.Module):
6
+ def __init__(self, num_classes):
7
+ super().__init__()
8
+ self.backbone = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
9
+
10
+ num_features = self.backbone.fc.in_features
11
+ self.backbone.fc = nn.Sequential(
12
+ nn.Dropout(0.5),
13
+ nn.Linear(num_features, 512),
14
+ nn.ReLU(),
15
+ nn.Dropout(0.3),
16
+ nn.Linear(512, num_classes)
17
+ )
18
+
19
+ def forward(self, x):
20
+ return self.backbone(x)