-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmodel.py
55 lines (50 loc) · 1.74 KB
/
model.py
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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
class Autoencoder(nn.Module):
def __init__(self):
super(Autoencoder, self).__init__()
# encoder layers
self.enc1 = nn.Conv1d(in_channels=8, out_channels=64, kernel_size=3, padding=1)
self.enc2 = nn.Conv1d(64, 32, kernel_size=3, padding=1)
self.enc3 = nn.Conv1d(32, 16, kernel_size=3, padding=1)
self.enc4 = nn.Conv1d(16, 8, kernel_size=3, padding=1)
self.pool = nn.MaxPool1d(2, 2)
# decoder layers
self.dec1 = nn.ConvTranspose1d(8, 8, kernel_size=2, stride=2)
self.dec2 = nn.ConvTranspose1d(8, 16, kernel_size=3, stride=2)
self.dec3 = nn.ConvTranspose1d(16, 32, kernel_size=2, stride=2)
self.dec4 = nn.ConvTranspose1d(32, 64, kernel_size=2, stride=2)
self.out = nn.Conv1d(64, 8, kernel_size=3)
def forward(self, x):
# encoder
# print(x.shape)
x = F.relu(self.enc1(x))
# print(x.shape)
x = self.pool(x)
# print(x.shape)
x = F.relu(self.enc2(x))
# print(x.shape)
x = self.pool(x)
# print(x.shape)
x = F.relu(self.enc3(x))
# print(x.shape)
x = self.pool(x)
# print(x.shape)
x = F.relu(self.enc4(x))
# print(x.shape)
x = self.pool(x) # the latent space representation
# print(x.shape)
# decoder
x = F.relu(self.dec1(x))
# print(x.shape)
x = F.relu(self.dec2(x))
# print(x.shape)
x = F.relu(self.dec3(x))
# print(x.shape)
x = F.relu(self.dec4(x))
# print(x.shape)
x = torch.sigmoid(self.out(x))
# print(x.shape)
return x