-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
88 lines (67 loc) · 3.32 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
84
85
86
87
88
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__()
# use for Xception model
self.conv1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, bias=False)
self.batch_norm1 = nn.BatchNorm2d(8, momentum=0.99, eps=1e-3)
self.relu1 = nn.ReLU()
self.conv2 = nn.Conv2d(in_channels=8, out_channels=8, kernel_size=3, bias=False)
self.batch_norm2 = nn.BatchNorm2d(8, momentum=0.99, eps=1e-3)
self.relu2 = nn.ReLU()
self.resblock1 = ResidualBlock(in_channeld=8, out_channels=16)
self.resblock2 = ResidualBlock(in_channeld=16, out_channels=32)
self.resblock3 = ResidualBlock(in_channeld=32, out_channels=64)
self.resblock4 = ResidualBlock(in_channeld=64, out_channels=128)
self.conv3 = nn.Conv2d(in_channels=128, out_channels=num_classes, kernel_size=3, padding=1, bias=False)
self.adapt_pool = nn.AdaptiveAvgPool2d(1)
def forward(self, input):
# use for Xception model
input = self.conv1(input)
input = self.batch_norm1(input)
input = self.relu1(input)
input = self.conv2(input)
input = self.batch_norm2(input)
input = self.relu2(input)
input = self.resblock1(input)
input = self.resblock2(input)
input = self.resblock3(input)
input = self.resblock4(input)
input = self.conv3(input)
input = self.adapt_pool(input)
input = input.view(input.size(0), -1)
return input