Nvd commited on
Commit
25cc53b
·
1 Parent(s): 1e98bba

Upload model.py

Browse files
Files changed (1) hide show
  1. model.py +34 -0
model.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch.nn as nn
2
+
3
+ class SimpleCNN(nn.Module):
4
+ def __init__(self):
5
+ super(SimpleCNN, self).__init__()
6
+ self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
7
+ self.relu1 = nn.ReLU()
8
+ self.pool1 = nn.MaxPool2d(2, 2)
9
+
10
+ self.conv2 = nn.Conv2d(16, 16, kernel_size=3, padding=1)
11
+ self.relu2 = nn.ReLU()
12
+ self.pool2 = nn.MaxPool2d(2, 2)
13
+
14
+ self.conv3 = nn.Conv2d(16, 16, kernel_size=3, padding=1)
15
+ self.relu3 = nn.ReLU()
16
+ self.pool3 = nn.MaxPool2d(2, 2)
17
+
18
+ self.conv4 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
19
+ self.relu4 = nn.ReLU()
20
+ self.pool4 = nn.MaxPool2d(2, 2)
21
+
22
+ self.fc1 = nn.Linear(32 * 2 * 2, 256)
23
+ self.fc2 = nn.Linear(256, 10)
24
+
25
+ def forward(self, x):
26
+ x = self.pool1(self.relu1(self.conv1(x)))
27
+ x = self.pool2(self.relu2(self.conv2(x)))
28
+ x = self.pool3(self.relu3(self.conv3(x)))
29
+ x = self.pool4(self.relu4(self.conv4(x)))
30
+ x = x.view(-1, 32 * 2 * 2)
31
+ x = self.relu4(self.fc1(x))
32
+ x = self.fc2(x)
33
+
34
+ return x