-
Notifications
You must be signed in to change notification settings - Fork 1
/
net_arch.py
183 lines (159 loc) · 7.72 KB
/
net_arch.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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
import torch
def make_vgg19_block(block):
"""Builds a vgg19 block from a dictionary
Args:
block: a dictionary
"""
layers = []
for i in range(len(block)):
one_ = block[i]
for k, v in one_.items():
if 'pool' in k:
layers += [torch.nn.MaxPool2d(kernel_size=v[0], stride=v[1],
padding=v[2])]
else:
conv2d = torch.nn.Conv2d(in_channels=v[0], out_channels=v[1],
kernel_size=v[2], stride=v[3],
padding=v[4])
layers += [conv2d, torch.nn.ReLU(inplace=True)]
return torch.nn.Sequential(*layers)
def make_stages(cfg_dict):
"""Builds CPM stages from a dictionary
Args:
cfg_dict: a dictionary
"""
stage_layer = torch.nn.ModuleList()
layers = []
for i in range(len(cfg_dict) - 1):
one_ = cfg_dict[i]
for k, v in one_.items():
if 'pool' in k:
layers += [torch.nn.MaxPool2d(kernel_size=v[0], stride=v[1],
padding=v[2])]
else:
conv2d = torch.nn.Conv2d(in_channels=v[0], out_channels=v[1],
kernel_size=v[2], stride=v[3],
padding=v[4])
stage_layer.append(conv2d)
stage_layer.append(torch.nn.ReLU(inplace=True))
one_ = list(cfg_dict[-1].keys())
k = one_[0]
v = cfg_dict[-1][k]
conv2d = torch.nn.Conv2d(in_channels=v[0], out_channels=v[1],
kernel_size=v[2], stride=v[3], padding=v[4])
stage_layer.append(conv2d)
return stage_layer # 注意这里不能用Sequential，因为要写densenet
def openpose_pami(stage_define = "PPPHHH"):
from config import generate_codec
stage_codec = generate_codec(stage_define)
num_stages = len(stage_define)
blocks = dict()
# block0是用了VGG前10层作为特征提取，这部分需要载入ImageNet权重
block0 = [{'conv1_1': [3, 64, 3, 1, 1]},
{'conv1_2': [64, 64, 3, 1, 1]},
{'pool1': [2, 2, 0]},
{'conv2_1': [64, 128, 3, 1, 1]},
{'conv2_2': [128, 128, 3, 1, 1]},
{'pool2': [2, 2, 0]},
{'conv3_1': [128, 256, 3, 1, 1]},
{'conv3_2': [256, 256, 3, 1, 1]},
{'conv3_3': [256, 256, 3, 1, 1]},
{'conv3_4': [256, 256, 3, 1, 1]},
{'pool3': [2, 2, 0]},
{'conv4_1': [256, 512, 3, 1, 1]},
{'conv4_2': [512, 512, 3, 1, 1]},
{'conv4_3': [512, 256, 3, 1, 1]},
{'conv4_4': [256, 128, 3, 1, 1]}]
for i in range(1, num_stages + 1):
# 输入这个stage的channel
if i==1:
# 上一个stage是VGG的输出
num_channel_in = 128
else:
# 上一个stage的输出和VGG的concate
num_channel_in = 128 + 19 * (stage_codec[i - 2] + 1)
# 输出channel，heatmap是19，PAF场是38
num_channel_out = 19 * (stage_codec[i - 1] + 1)
# 制作当前stage的网络结构
blocks['block%d' % i] = [
{'conv_stage%d_denseblock1_1' % i: [num_channel_in, 128, 3, 1, 1]},
{'conv_stage%d_denseblock1_2' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock1_3' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock2_1' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock2_2' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock2_3' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock3_1' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock3_2' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock3_3' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock4_1' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock4_2' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock4_3' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock5_1' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock5_2' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_denseblock5_3' % i: [128, 128, 3, 1, 1]},
{'conv_stage%d_6' % i: [128, 128, 1, 1, 0]},
{'conv_stage%d_7' % i: [128, num_channel_out, 1, 1, 0]}
]
# 构造模型字典
models = dict()
models['block0'] = make_vgg19_block(block0)
for k, v in blocks.items():
models[k] = make_stages(list(v))
# 构造网络模型
class cmu_openpose_model(torch.nn.Module):
def __init__(self, model_dict):
super(cmu_openpose_model, self).__init__()
self.model0 = model_dict['block0']
# 采用动态执行代码的方式，构建任意stage的网络模型
for i in range(1, num_stages + 1):
exec("self.model%d = model_dict['block%d']"%(i, i))
# 随机初始化
self._initialize_weights_norm()
def forward(self, x):
# 定义intermediate_map用来存放每个stage的map信息
intermediate_map = []
# out0用来存储VGG19的前10层信息
out0 = self.model0(x)
# 实现stages
for i_stage in range(1, num_stages + 1):
if i_stage > 1:
# 除了第一个stage以外，每个stage使用上一个stage的结果加上VGG feature
out1 = torch.cat([out0, out1], 1)
else:
# 第一个stage只用VGG feature
out1 = out0
# 对连续5个块的densenet的实现
for i in range(5):
out1_d1 = eval("self.model%d[i*6+0](out1)" % i_stage)
out1_d1 = eval("self.model%d[i*6+1](out1_d1)" % i_stage)
out1_d2 = eval("self.model%d[i*6+2](out1_d1)" % i_stage)
out1_d2 = eval("self.model%d[i*6+3](out1_d2)" % i_stage)
out1_d3 = eval("self.model%d[i*6+4](out1_d2)" % i_stage)
out1_d3 = eval("self.model%d[i*6+5](out1_d3)" % i_stage)
out1 = out1_d1 + out1_d2 + out1_d3
# 然后是2个1*1整合
out1 = eval("self.model%d[31](self.model%d[30](out1))" % (i_stage, i_stage))
out1 = eval("self.model%d[32](out1)" % i_stage)
# 这里out1得到的就是PAF或者confidence map
# 记录下这个map用于intermediate supervision
intermediate_map.append(out1)
# 每个stage的实现到此结束
# 不能在exec函数中调用return，否则报错"SyntaxError: 'return' outside function"
# exec("return (out%d_1, out%d_2), saved_for_loss"%(num_stages, num_stages))
# 因此改用eval动态赋值的方式
#(eval("out%d_1"%num_stages), eval("out%d_2"%num_stages)), saved_for_loss
return intermediate_map
def _initialize_weights_norm(self):
for m in self.modules():
if isinstance(m, torch.nn.Conv2d):
torch.nn.init.normal_(m.weight, std=0.01)
# 预留给mobilenet模型用(如果后面想起来写的话)，因为他的conv2d没有bias
if m.bias is not None:
torch.nn.init.constant_(m.bias, 0.0)
# 每个stage最后一层没有Relu
torch.nn.init.normal_(self.model1[32].weight, std=0.01)
for i in range(2, 4):
exec("torch.nn.init.normal_(self.model%d[32].weight, std=0.01)"%i)
model = cmu_openpose_model(models)
print("Openpose model building finished.")
return model