-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathmodel.py
More file actions
32 lines (22 loc) · 986 Bytes
/
model.py
File metadata and controls
32 lines (22 loc) · 986 Bytes
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import torch
import torch.nn as nn
# Main CNN model
class Model(nn.Module):
def __init__(self, n):
super(Model, self).__init__()
self.relu = nn.ReLU()
self.tanh = nn.Tanh()
self.conv3_32 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
self.conv32_32 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
self.conv64_32 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1)
self.conv64_3n = nn.Conv2d(64, 3*n, kernel_size=3, stride=1, padding=1)
# Total of 7 layers with skip connections
def forward(self, inpt):
output1 = self.relu(self.conv3_32(inpt.float()))
output2 = self.relu(self.conv32_32(output1))
output3 = self.relu(self.conv32_32(output2))
output4 = self.relu(self.conv32_32(output3))
output5 = self.relu(self.conv64_32(torch.cat([output4, output3], dim=1)))
output6 = self.relu(self.conv64_32(torch.cat([output5, output2], dim=1)))
output7 = self.tanh(self.conv64_3n(torch.cat([output6, output1], dim=1)))
return output7