FisayoF commited on
Commit
67f122a
·
verified ·
1 Parent(s): b016e9e

Add model definition

Browse files
Files changed (1) hide show
  1. model.py +26 -0
model.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+
4
+ class CnnNet(nn.Module):
5
+ def __init__(self):
6
+ super(CnnNet, self).__init__()
7
+ # Spectrogram is single channel, so use 1 input rather than 3 (RGB)
8
+ self.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1)
9
+ self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
10
+ self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
11
+
12
+ self.fc1 = nn.Linear(128 * 120 * 160, 3)
13
+
14
+ def forward(self, x):
15
+ x = self.conv1(x)
16
+ x = F.relu(x)
17
+ x = self.pool(x)
18
+
19
+ x = self.conv2(x)
20
+ x = F.relu(x)
21
+ x = self.pool(x)
22
+
23
+ x = x.view(-1, 128 * 120 * 160)
24
+ x = self.fc1(x)
25
+ x = F.log_softmax(x, dim=1)
26
+ return x