blatoet commited on
Commit
beda2dd
·
verified ·
1 Parent(s): 2f63ee2

Upload model.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. model.py +34 -0
model.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+
4
+ class CatDogClassifier(nn.Module):
5
+ def __init__(self):
6
+ super(CatDogClassifier, self).__init__()
7
+
8
+ # Convolutional layers
9
+ self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
10
+ self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
11
+ self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
12
+
13
+ # Pooling and activation
14
+ self.pool = nn.MaxPool2d(2, 2)
15
+ self.relu = nn.ReLU()
16
+ self.dropout = nn.Dropout(0.5)
17
+
18
+ # Fully connected layers
19
+ self.fc1 = nn.Linear(128 * 28 * 28, 512)
20
+ self.fc2 = nn.Linear(512, 128)
21
+ self.fc3 = nn.Linear(128, 2)
22
+
23
+ def forward(self, x):
24
+ x = self.relu(self.conv1(x))
25
+ x = self.pool(x)
26
+ x = self.relu(self.conv2(x))
27
+ x = self.pool(x)
28
+ x = self.relu(self.conv3(x))
29
+ x = self.pool(x)
30
+ x = x.view(x.size(0), -1)
31
+ x = self.dropout(self.relu(self.fc1(x)))
32
+ x = self.dropout(self.relu(self.fc2(x)))
33
+ x = self.fc3(x)
34
+ return x