-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
83 lines (66 loc) · 3.18 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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import torch.nn as nn
class SeparableConv2d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False):
super(SeparableConv2d, self).__init__()
self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size, stride, padding, dilation, groups=in_channels,
bias=bias)
self.pointwise = nn.Conv2d(in_channels, out_channels, 1, 1, 0, 1, 1, bias=bias)
def forward(self, x):
x = self.depthwise(x)
x = self.pointwise(x)
return x
class ResidualBlock(nn.Module):
def __init__(self, in_channeld, out_channels):
super(ResidualBlock, self).__init__()
self.residual_conv = nn.Conv2d(in_channels=in_channeld, out_channels=out_channels, kernel_size=1, stride=2,
bias=False)
self.residual_bn = nn.BatchNorm2d(out_channels, momentum=0.99, eps=1e-3)
self.sepConv1 = SeparableConv2d(in_channels=in_channeld, out_channels=out_channels, kernel_size=3, bias=False,
padding=1)
self.bn1 = nn.BatchNorm2d(out_channels, momentum=0.99, eps=1e-3)
self.relu = nn.ReLU()
self.sepConv2 = SeparableConv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, bias=False,
padding=1)
self.bn2 = nn.BatchNorm2d(out_channels, momentum=0.99, eps=1e-3)
self.maxp = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
def forward(self, x):
res = self.residual_conv(x)
res = self.residual_bn(res)
x = self.sepConv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.sepConv2(x)
x = self.bn2(x)
x = self.maxp(x)
return res + x
class Model(nn.Module):
def __init__(self, num_classes):
super(Model, self).__init__()
self.conv1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, stride=1, bias=False)
self.bn1 = nn.BatchNorm2d(8, affine=True, momentum=0.99, eps=1e-3)
self.relu1 = nn.ReLU()
self.conv2 = nn.Conv2d(in_channels=8, out_channels=8, kernel_size=3, stride=1, bias=False)
self.bn2 = nn.BatchNorm2d(8, momentum=0.99, eps=1e-3)
self.relu2 = nn.ReLU()
self.module1 = ResidualBlock(in_channeld=8, out_channels=16)
self.module2 = ResidualBlock(in_channeld=16, out_channels=32)
self.module3 = ResidualBlock(in_channeld=32, out_channels=64)
self.module4 = ResidualBlock(in_channeld=64, out_channels=128)
self.last_conv = nn.Conv2d(in_channels=128, out_channels=num_classes, kernel_size=3, padding=1)
self.avgp = nn.AdaptiveAvgPool2d((1, 1))
def forward(self, input):
x = input
x = self.conv1(x)
x = self.bn1(x)
x = self.relu1(x)
x = self.conv2(x)
x = self.bn2(x)
x = self.relu2(x)
x = self.module1(x)
x = self.module2(x)
x = self.module3(x)
x = self.module4(x)
x = self.last_conv(x)
x = self.avgp(x)
x = x.view((x.shape[0], -1))
return x