-
Notifications
You must be signed in to change notification settings - Fork 1
/
FKPStructure.py
36 lines (34 loc) · 1.4 KB
/
FKPStructure.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
from torch.nn import Linear, ReLU, Sequential, Conv2d, MaxPool2d, Module, BatchNorm2d, AvgPool2d
class FKPStructure(Module):
def __init__(self):
super(FKPStructure, self).__init__()
#Have 3 Covolution layer for feature extraction
self.cnn_layers = Sequential(
#first Convolution layer and Max pooling
Conv2d(1, 4, kernel_size=3, stride=1, padding=1, padding_mode='reflect'),
BatchNorm2d(4),
ReLU(inplace=True),
MaxPool2d(kernel_size=2, stride=2),
#second Convolution layer and Max pooling
Conv2d(4, 8, kernel_size=3, stride=1, padding=1, padding_mode='reflect'),
BatchNorm2d(8),
ReLU(inplace=True),
MaxPool2d(kernel_size=2, stride=2),
#third Convolution layer and Average pooling
Conv2d(8, 16, kernel_size=3, stride=1, padding=1, padding_mode='reflect'),
BatchNorm2d(16),
ReLU(inplace=True),
AvgPool2d(kernel_size=2, stride=2),
#output is (16 *13.5 * 27.5)
)
#Have 1 Fully Connected layer for classification
self.linear_layers1 = Sequential(
Linear(16 *13 * 27, 100)
)
#data flow
def forward(self, x):
x = self.cnn_layers(x)
# convert feature map to vector form
x = x.view(x.size(0), -1)
x = self.linear_layers1(x)
return x