-
Notifications
You must be signed in to change notification settings - Fork 2
/
Resnet.py
93 lines (77 loc) · 3.7 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
from Config import opt
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_blocks(nn.Module):
def __init__(self, num_layer_stack, img_chns, img_res):
super(ResNet_blocks, self).__init__()
self.conv1 = _conv2d_bn_relu(in_channels=img_chns, 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.AdaptiveAvgPool2d((1, 1))
self.dz = 64
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)
gap = self.avgpool(y)
gap = gap.view(gap.size(0), -1)
return gap
class ResNet_cf(nn.Module):
def __init__(self, label_num):
super(ResNet_cf, self).__init__()
self.fc = nn.Linear(in_features=64, out_features=label_num)
def forward(self, gap):
# y = gap.view(gap.size(0), -1)
y = self.fc(gap)
return y
class ResNet(nn.Module):
def __init__(self, num_layer_stack, label_num, img_chns, img_res, dataset='mnist', train_num=10):
super(ResNet, self).__init__()
self.ResNet_block = ResNet_blocks(num_layer_stack, img_chns, img_res)
self.ResNet_cf = ResNet_cf(label_num)
def forward(self, img):
gap = self.ResNet_block(img)
y = self.ResNet_cf(gap)
return y, gap