a b/py_version/model.py
1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4
5
class ConvBlock(nn.Module):
6
    
7
    def __init__(self, in_channels, out_channels, kernel, stride):
8
        super(ConvBlock, self).__init__()
9
10
        self.layer = nn.Sequential(
11
                    nn.Conv1d(in_channels = in_channels, out_channels = out_channels, kernel_size = kernel, stride = stride),
12
                    nn.BatchNorm1d(out_channels),
13
                    nn.ReLU()
14
                    )
15
16
    def forward(self, x):
17
18
        out = self.layer(x)
19
20
        return out
21
22
23
class CNN1D(nn.Module):
24
    
25
    def __init__(self, num_classes):
26
        super(CNN1D, self).__init__()
27
28
        self.layer = nn.Sequential(
29
                    ConvBlock(16, 16, 10, 4),
30
                    ConvBlock(16, 16, 5, 2),
31
                    ConvBlock(16, 16, 5, 2),
32
                    ConvBlock(16, 32, 5, 2),
33
                    nn.MaxPool1d(kernel_size = 2, stride = 2),
34
                    ConvBlock(32, 32, 4, 1),
35
                    ConvBlock(32, 32, 4, 1),
36
                    ConvBlock(32, 64, 4, 1),
37
                    nn.MaxPool1d(kernel_size = 2, stride = 2),
38
                    ConvBlock(64, 64, 3, 1),
39
                    ConvBlock(64, 64, 3, 1),
40
                    ConvBlock(64, 128, 3, 1),
41
                    nn.MaxPool1d(kernel_size = 2, stride = 2),
42
                    ConvBlock(128, 128, 2, 1),
43
                    ConvBlock(128, 128, 2, 1),
44
                    ConvBlock(128, 256, 2, 1),
45
                    nn.MaxPool1d(kernel_size = 2, stride = 2)
46
                    )
47
48
        self.linear = nn.Sequential(
49
                    nn.Linear(1280, 512),
50
                    nn.ReLU(),
51
                    nn.Dropout(0.5),
52
                    nn.Linear(512, 128),
53
                    nn.ReLU(),
54
                    nn.Linear(128, num_classes)
55
                    )
56
57
    def forward(self, x):
58
59
        batch_size = x.shape[0]
60
        out = self.layer(x)
61
        out = out.view(batch_size, -1)
62
        out = self.linear(out)
63
64
        return out
65
66
67
class CNN1D_F(nn.Module):
68
    
69
    def __init__(self, num_classes):
70
        super(CNN1D_F, self).__init__()
71
72
        self.layer = nn.Sequential(
73
                    ConvBlock(16, 16, 10, 4),
74
                    ConvBlock(16, 16, 5, 2),
75
                    ConvBlock(16, 16, 5, 2),
76
                    nn.MaxPool1d(kernel_size = 2, stride = 2),
77
                    ConvBlock(16, 16, 5, 2),
78
                    ConvBlock(16, 16, 5, 2),
79
                    ConvBlock(16, 32, 5, 2),
80
                    nn.MaxPool1d(kernel_size = 2, stride = 2),
81
                    ConvBlock(32, 32, 4, 1),
82
                    ConvBlock(32, 32, 4, 1),
83
                    ConvBlock(32, 64, 4, 1),
84
                    nn.MaxPool1d(kernel_size = 2, stride = 2),
85
                    ConvBlock(64, 64, 3, 1),
86
                    ConvBlock(64, 64, 3, 1),
87
                    ConvBlock(64, 128, 3, 1),
88
                    nn.MaxPool1d(kernel_size = 2, stride = 2),
89
                    ConvBlock(128, 128, 2, 1),
90
                    ConvBlock(128, 128, 2, 1),
91
                    ConvBlock(128, 256, 2, 1),
92
                    nn.MaxPool1d(kernel_size = 2, stride = 2)
93
                    )
94
95
        self.linear = nn.Sequential(
96
                    nn.Linear(1280, 512),
97
                    nn.ReLU(),
98
                    nn.Dropout(0.5),
99
                    nn.Linear(512, 128),
100
                    nn.ReLU(),
101
                    nn.Linear(128, num_classes)
102
                    )
103
104
    def forward(self, x):
105
106
        batch_size = x.shape[0]
107
        out = self.layer(x)
108
        out = out.view(batch_size, -1)
109
        out = self.linear(out)
110
111
        return out