/
resnet.py
91 lines (73 loc) · 3.61 KB
/
resnet.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
89
90
91
import torch.nn as nn
import math
def _conv2d_bn(in_channels, out_channels, kernel_size, stride, padding):
conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
kernel_size=kernel_size, stride=stride, padding=padding, bias=False)
bn = nn.BatchNorm2d(num_features=out_channels)
return nn.Sequential(conv, bn)
def _conv2d_bn_relu(in_channels, out_channels, kernel_size, stride, padding):
conv2d_bn = _conv2d_bn(in_channels, out_channels, kernel_size, stride, padding)
relu = nn.ReLU(inplace=True)
layers = list(conv2d_bn.children())
layers.append(relu)
return nn.Sequential(*layers)
class _BasicBlock(nn.Module):
def __init__(self, in_channels, out_channels, downscale=False):
super(_BasicBlock, self).__init__()
self.down_sampler = None
stride = 1
if downscale:
self.down_sampler = _conv2d_bn(in_channels, out_channels, kernel_size=1, stride=2, padding=0)
stride = 2
self.conv_bn_relu1 = _conv2d_bn_relu(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
# don't relu here! relu on (H(x) + x)
self.conv_bn2 = _conv2d_bn(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
self.relu_out = nn.ReLU(inplace=True)
# residual = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
# residual = nn.BatchNorm2d(num_features=out_channels)
# residual = nn.ReLU(inplace=True)
def forward(self, x):
input = x
if self.down_sampler:
input = self.down_sampler(x)
residual = self.conv_bn_relu1(x)
residual = self.conv_bn2(residual)
out = self.relu_out(input + residual)
return out
class _ResNet(nn.Module):
def __init__(self, num_layer_stack):
super(_ResNet, self).__init__()
self.conv1 = _conv2d_bn_relu(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
self.layer1 = self.__make_layers(num_layer_stack, in_channels=16, out_channels=16, downscale=False)
self.layer2 = self.__make_layers(num_layer_stack, in_channels=16, out_channels=32, downscale=True)
self.layer3 = self.__make_layers(num_layer_stack, in_channels=32, out_channels=64, downscale=True)
self.avgpool = nn.AvgPool2d(kernel_size=8, stride=1)
self.fc = nn.Linear(in_features=64, out_features=10)
# for m in self.modules():
# if isinstance(m, nn.Conv2d):
# n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
# m.weight.data.normal_(0, math.sqrt(2. / n))
# elif isinstance(m, nn.BatchNorm2d):
# m.weight.data.fill_(1)
# m.bias.data.zero_()
def __make_layers(self, num_layer_stack, in_channels, out_channels, downscale):
layers = []
layers.append(_BasicBlock(in_channels=in_channels, out_channels=out_channels, downscale=downscale))
for i in range(num_layer_stack - 1):
layers.append(_BasicBlock(in_channels=out_channels, out_channels=out_channels, downscale=False))
return nn.Sequential(*layers)
def forward(self, x):
y = self.conv1(x)
y = self.layer1(y)
y = self.layer2(y)
y = self.layer3(y)
y = self.avgpool(y)
y = y.view(y.size(0), -1)
y = self.fc(y)
return y
def resnet20():
return _ResNet(num_layer_stack=3)
def resnet32():
return _ResNet(num_layer_stack=5)
def resnet56():
return _ResNet(num_layer_stack=9)