1#---------------------------------------------------------------------------------------------- 2# Copyright (c) Microsoft Corporation. All rights reserved. 3# Licensed under the MIT License. See License.txt in the project root for license information. 4#---------------------------------------------------------------------------------------------- 5 6from __future__ import absolute_import 7from __future__ import division 8from __future__ import print_function 9 10import numpy as np 11from collections import OrderedDict 12from mmdnn.conversion.common.DataStructure.graph import GraphNode, Graph 13# from tensorflow.core.framework.node_def_pb2 import NodeDef 14# from tensorflow.core.framework import attr_value_pb2 15 16 17class DarknetGraphNode(GraphNode): 18 19 def __init__(self, layer): 20 21 super(DarknetGraphNode, self).__init__(layer) 22 23 24 @property 25 def name(self): 26 return self.layer['name'] 27 28 29 @property 30 def type(self): 31 return self.layer['type'] 32 33 34 @property 35 def dk_layer(self): 36 return self.layer 37 38 39 def get_attr(self, name, default_value = None): 40 if name in self.layer['attr'].keys(): 41 return self.layer['attr'][name] 42 else: 43 return default_value 44 45 46class DarknetGraph(Graph): 47 48 def __init__(self, model): 49 # pass 50 51 super(DarknetGraph, self).__init__(model) 52 self.layer_num_map = {} 53 self.model = model 54 self.weights = {} 55 self.original_list = OrderedDict() 56 57 @staticmethod 58 def dim_str_to_int(input_dim): 59 if type(input_dim) == list: 60 return [int(i) for i in input_dim] 61 62 @staticmethod 63 def conv_output_width(width, padding, kernel_size, stride): 64 return (width + 2*padding - kernel_size)/stride + 1 65 66 @staticmethod 67 def conv_output_height(height, padding, kernel_size, stride): 68 return (height + 2*padding - kernel_size)/stride + 1 69 70 def build(self): 71 72 for i, block in enumerate(self.model): 73 # print("\n") 74 # print(i) 75 # print(block) 76 77 # continue 78 node = OrderedDict() 79 if block['type'] == 'net': 80 node['name'] = 'dk_Input' 81 node['input'] = ['data'] 82 node['type'] = 'DataInput' 83 node['input_dim'] = ['-1'] 84 # NHWC 85 node['input_dim'].append(block['height']) 86 node['input_dim'].append(block['width']) 87 node['input_dim'].append(block['channels']) 88 input_param = OrderedDict() 89 input_param['shape'] = self.dim_str_to_int(node['input_dim']) 90 input_param['_output_shape'] = self.dim_str_to_int(node['input_dim']) 91 node['attr'] = input_param 92 self.layer_map[node['name']] = DarknetGraphNode(node) 93 self.original_list[node['name']] = DarknetGraphNode(node) 94 self.layer_num_map[i] = node['name'] 95 pre_node_name = node['name'] 96 97 elif block['type'] == 'convolutional': 98 conv_layer = OrderedDict() 99 conv_layer['input'] = [pre_node_name] 100 101 input_shape = self.layer_map[pre_node_name].get_attr('_output_shape') 102 w = input_shape[1] 103 h = input_shape[2] 104 channels = input_shape[3] 105 # assert False 106 107 if 'name' in block.keys(): 108 conv_layer['name'] = block['name'] 109 else: 110 conv_layer['name'] = 'layer%d-conv' % i 111 conv_layer['type'] = 'Conv' 112 113 convolution_param = OrderedDict() 114 convolution_param['num_output'] = int(block['filters']) 115 convolution_param['kernel_size'] = int(block['size']) 116 convolution_param['kernel'] = [int(block['size']), int(block['size']), channels, int(block['filters'])] 117 convolution_param['pad'] = int(block['pad']) 118 119 if block['pad'] == '1': 120 convolution_param['padding'] = int(convolution_param['kernel_size'])//2 121 convolution_param['stride'] = int(block['stride']) 122 if block['batch_normalize'] == '1': 123 convolution_param['bias_term'] = 'false' 124 else: 125 convolution_param['bias_term'] = 'true' 126 output_w = self.conv_output_width(w ,convolution_param['padding'], convolution_param['kernel_size'], convolution_param['stride']) 127 output_h = self.conv_output_height(h ,convolution_param['padding'], convolution_param['kernel_size'], convolution_param['stride']) 128 convolution_param['_output_shape'] = [-1, output_w, output_h, convolution_param['num_output']] 129 conv_layer['attr'] = convolution_param 130 self.layer_map[conv_layer['name']] = DarknetGraphNode(conv_layer) 131 self.original_list[conv_layer['name']] = DarknetGraphNode(conv_layer) 132 pre_node_name = conv_layer['name'] 133 134 if block['batch_normalize'] == '1': 135 bn_layer = OrderedDict() 136 bn_layer['input'] = [pre_node_name] 137 138 139 input_shape = self.layer_map[pre_node_name].get_attr('_output_shape') 140 if 'name' in block.keys(): 141 bn_layer['name'] = '%s-bn' % block['name'] 142 else: 143 bn_layer['name'] = 'layer%d-bn' % i 144 bn_layer['type'] = 'BatchNorm' 145 batch_norm_param = OrderedDict() 146 batch_norm_param['use_global_stats'] = True 147 batch_norm_param['_output_shape'] = convolution_param['_output_shape'] 148 batch_norm_param['bias_term'] = True 149 batch_norm_param['scale'] = True 150 bn_layer['attr'] = batch_norm_param 151 152 153 self.layer_map[bn_layer['name']] = DarknetGraphNode(bn_layer) 154 self.original_list[bn_layer['name']] = DarknetGraphNode(bn_layer) 155 156 pre_node_name = bn_layer['name'] 157 158 159 if block['activation'] != 'linear': 160 relu_layer = OrderedDict() 161 relu_layer['input'] = [pre_node_name] 162 if 'name' in block.keys(): 163 relu_layer['name'] = '%s-act' % block['name'] 164 else: 165 relu_layer['name'] = 'layer%d-act' % i 166 relu_layer['type'] = 'ReLU' 167 relu_param = OrderedDict() 168 if block['activation'] == 'leaky': 169 relu_layer['type'] = 'leakyReLU' 170 relu_param['negative_slope'] = '0.1' 171 relu_param['_output_shape'] = input_shape 172 relu_layer['attr'] = relu_param 173 self.layer_map[relu_layer['name']] = DarknetGraphNode(relu_layer) 174 self.layer_num_map[i] = relu_layer['name'] 175 self.original_list[relu_layer['name']] = DarknetGraphNode(relu_layer) 176 pre_node_name = relu_layer['name'] 177 178 else: 179 self.layer_num_map[i] = bn_layer['name'] 180 181 182 elif block['type'] == 'maxpool': 183 max_layer = OrderedDict() 184 max_layer['input'] = [pre_node_name] 185 if 'name' in block.keys(): 186 max_layer['name'] = block['name'] 187 else: 188 max_layer['name'] = 'layer%d-maxpool' % i 189 max_layer['type'] = 'Pooling' 190 pooling_param = OrderedDict() 191 pooling_param['kernel_size'] = int(block['size']) 192 pooling_param['stride'] = int(block['stride']) 193 pooling_param['pool'] = 'MAX' 194 pooling_param['padding'] = 0 195 if 'pad' in block.keys() and int(block['pad']) == 1: 196 pooling_param['padding'] = (int(block['size'])-1)/2 197 198 input_shape = self.layer_map[pre_node_name].get_attr('_output_shape') 199 w = input_shape[1] 200 h = input_shape[2] 201 output_w = (w + 2*pooling_param['padding'])/pooling_param['stride'] 202 output_h = (h + 2*pooling_param['padding'])/pooling_param['stride'] 203 204 pooling_param['_output_shape'] = [-1, output_w, output_h, input_shape[-1]] 205 max_layer['attr'] = pooling_param 206 self.layer_map[max_layer['name']] = DarknetGraphNode(max_layer) 207 self.original_list[max_layer['name']] = DarknetGraphNode(max_layer) 208 self.layer_num_map[i] = max_layer['name'] 209 pre_node_name = max_layer['name'] 210 211 elif block['type'] == 'avgpool': 212 avg_layer = OrderedDict() 213 214 avg_layer['input'] = [pre_node_name] 215 if 'name' in block.keys(): 216 avg_layer['name'] = block['name'] 217 else: 218 avg_layer['name'] = 'layer%d-avgpool' % i 219 avg_layer['type'] = 'Pooling' 220 pooling_param = OrderedDict() 221 input_shape = self.layer_map[pre_node_name].get_attr('_output_shape') 222 pooling_param['_output_shape'] = [-1, 1, 1, input_shape[-1]] 223 pooling_param['pool'] = 'AVG' 224 avg_layer['attr'] = pooling_param 225 self.layer_map[avg_layer['name']] = DarknetGraphNode(avg_layer) 226 self.original_list[avg_layer['name']] = DarknetGraphNode(avg_layer) 227 self.layer_num_map[i] = avg_layer['name'] 228 pre_node_name = avg_layer['name'] 229 230 elif block['type'] == 'route': 231 prev = block['layers'].split(',') #[-1,61] 232 if len(prev) == 1: 233 prev_layer_id = i + int(prev[0]) 234 self.layer_num_map[i] = self.layer_num_map[prev_layer_id] 235 pre_node_name = self.layer_num_map[i] 236 elif len(prev) == 2: 237 input_list = [] 238 input_shape = [] 239 route_layer = OrderedDict() 240 for p in prev: 241 if int(p)>0: 242 243 input_name = self.layer_num_map[int(p)+1] 244 input_list.append(input_name) 245 input_shape.append(self.layer_map[input_name].get_attr('_output_shape')) 246 247 else: 248 prev_layer_id = i + int(p) 249 input_name = self.layer_num_map[prev_layer_id] 250 input_shape.append(self.layer_map[input_name].get_attr('_output_shape')) 251 input_list.append(input_name) 252 route_param = OrderedDict() 253 254 255 shape_ = 0 256 for shape in input_shape: 257 shape_ += shape[-1] 258 route_param['axis'] = 3 259 route_param['_output_shape'] = input_shape[0][:-1] + [shape_] 260 route_layer['input'] = input_list 261 262 if 'name' in block.keys(): 263 route_layer['name'] = block['name'] 264 else: 265 route_layer['name'] = 'layer%d-concat' % i 266 267 route_layer['type'] = 'Concat' 268 route_layer['attr'] = route_param 269 270 self.layer_map[route_layer['name']] = DarknetGraphNode(route_layer) 271 self.original_list[route_layer['name']] = DarknetGraphNode(route_layer) 272 self.layer_num_map[i] = route_layer['name'] 273 pre_node_name = route_layer['name'] 274 275 elif block['type'] == 'shortcut': 276 prev_layer_id1 = i + int(block['from']) 277 prev_layer_id2 = i - 1 278 bottom1 = self.layer_num_map[prev_layer_id1] 279 bottom2 = self.layer_num_map[prev_layer_id2] 280 input_shape = self.layer_map[bottom2].get_attr('_output_shape') 281 shortcut_layer = OrderedDict() 282 shortcut_layer['input'] = [bottom1, bottom2] 283 # print(shortcut_layer['input'] ) 284 if 'name' in block.keys(): 285 shortcut_layer['name'] = block['name'] 286 else: 287 shortcut_layer['name'] = 'layer%d-shortcut' % i 288 shortcut_layer['type'] = 'Add' 289 eltwise_param = OrderedDict() 290 eltwise_param['operation'] = 'SUM' 291 eltwise_param['_output_shape'] = input_shape 292 shortcut_layer['attr'] = eltwise_param 293 294 295 self.layer_map[shortcut_layer['name']] = DarknetGraphNode(shortcut_layer) 296 self.original_list[shortcut_layer['name']] = DarknetGraphNode(shortcut_layer) 297 self.layer_num_map[i] = shortcut_layer['name'] 298 pre_node_name = shortcut_layer['name'] 299 300 if block['activation'] != 'linear': 301 relu_layer = OrderedDict() 302 relu_layer['input'] = [pre_node_name] 303 if 'name' in block.keys(): 304 relu_layer['name'] = '%s-act' % block['name'] 305 else: 306 relu_layer['name'] = 'layer%d-act' % i 307 relu_layer['type'] = 'ReLU' 308 relu_param = OrderedDict() 309 relu_param['_output_shape'] = input_shape 310 if block['activation'] == 'leaky': 311 312 relu_param['negative_slope'] = '0.1' 313 314 relu_layer['attr'] = relu_param 315 self.layer_map[relu_layer['name']] = DarknetGraphNode(relu_layer) 316 self.original_list[relu_layer['name']] = DarknetGraphNode(relu_layer) 317 pre_node_name = relu_layer['name'] 318 319 elif block['type'] == 'connected': 320 fc_layer = OrderedDict() 321 fc_layer['input'] = [pre_node_name] 322 if 'name' in block.keys(): 323 fc_layer['name'] = block['name'] 324 else: 325 fc_layer['name'] = 'layer%d-fc' % i 326 fc_layer['type'] = 'InnerProduct' 327 fc_param = OrderedDict() 328 fc_param['num_output'] = int(block['output']) 329 input_shape = self.layer_map[pre_node_name].get_attr('_output_shape') 330 fc_param['_output_shape'] = input_shape[:-1] + [fc_param['num_output']] 331 fc_layer['attr'] = fc_param 332 self.layer_map[fc_layer['name']] = DarknetGraphNode(fc_layer) 333 self.original_list[fc_layer['name']] = DarknetGraphNode(fc_layer) 334 self.layer_num_map[i] = fc_layer['name'] 335 pre_node_name = fc_layer['name'] 336 337 if block['activation'] != 'linear': 338 relu_layer = OrderedDict() 339 relu_layer['input'] = [pre_node_name] 340 if 'name' in block.keys(): 341 relu_layer['name'] = '%s-act' % block['name'] 342 else: 343 relu_layer['name'] = 'layer%d-act' % i 344 relu_layer['type'] = 'ReLU' 345 relu_param = OrderedDict() 346 if block['activation'] == 'leaky': 347 348 relu_param['negative_slope'] = '0.1' 349 relu_param['_output_shape'] = fc_param['_output_shape'] 350 relu_layer['attr'] = relu_param 351 self.layer_map[relu_layer['name']] = DarknetGraphNode(relu_layer) 352 self.original_list[relu_layer['name']] = DarknetGraphNode(relu_layer) 353 pre_node_name = relu_layer['name'] 354 355 elif block['type'] == 'softmax': 356 sm_layer = OrderedDict() 357 358 sm_layer['input'] = [pre_node_name] 359 if 'name' in block.keys(): 360 sm_layer['name'] = block['name'] 361 else: 362 sm_layer['name'] = 'layer%d-softmax' % i 363 sm_layer['type'] = 'Softmax' 364 softmax_param = OrderedDict() 365 input_shape = self.layer_map[pre_node_name].get_attr('_output_shape') 366 softmax_param['_output_shape'] = input_shape 367 sm_layer['attr'] = softmax_param 368 self.layer_map[sm_layer['name']] = DarknetGraphNode(sm_layer) 369 self.original_list[sm_layer['name']] = DarknetGraphNode(sm_layer) 370 self.layer_num_map[i] = sm_layer['name'] 371 pre_node_name = sm_layer['name'] 372 373 elif block['type'] == 'yolo': 374 375 yolo_layer = OrderedDict() 376 yolo_layer['input'] = [pre_node_name] 377 if 'name' in block.keys(): 378 yolo_layer['name'] = block['name'] 379 else: 380 yolo_layer['name'] = 'layer%d-yolo' % i 381 yolo_layer['type'] = 'yolo' 382 yolo_param = OrderedDict() 383 yolo_param['truth_thresh'] = float(block['truth_thresh']) 384 yolo_param['random'] = float(block['random']) 385 yolo_param['ignore_thresh'] = float(block['ignore_thresh']) 386 yolo_param['jitter'] = float(block['jitter']) 387 yolo_param['num'] = int(block['num']) 388 yolo_param['classes'] = int(block['classes']) 389 anchors = [int(t) for t in block['anchors'].split(',')] 390 yolo_param['anchors'] = anchors 391 mask = [int(t) for t in block['mask'].split(',')] 392 yolo_param['mask'] = mask 393 394 yolo_layer['attr'] = yolo_param 395 self.layer_map[yolo_layer['name']] = DarknetGraphNode(yolo_layer) 396 self.original_list[yolo_layer['name']] = DarknetGraphNode(yolo_layer) 397 self.layer_num_map[i] = yolo_layer['name'] 398 399 elif block['type'] == 'upsample': 400 401 input_shape = self.layer_map[pre_node_name].get_attr('_output_shape') 402 upsample_layer = OrderedDict() 403 upsample_layer['input'] = [pre_node_name] 404 if 'name' in block.keys(): 405 upsample_layer['name'] = block['name'] 406 else: 407 upsample_layer['name'] = 'layer%d-upsample' % i 408 upsample_layer['type'] = 'upsample' 409 upsample_param = OrderedDict() 410 stride = block['stride'] 411 upsample_param['scales'] = [int(stride), int(stride)] 412 upsample_param['_output_shape'] = [input_shape[0]] + [q*int(stride) for q in input_shape[1:3]] + [input_shape[-1]] 413 upsample_layer['attr'] = upsample_param 414 self.layer_map[upsample_layer['name']] = DarknetGraphNode(upsample_layer) 415 self.original_list[upsample_layer['name']] = DarknetGraphNode(upsample_layer) 416 self.layer_num_map[i] = upsample_layer['name'] 417 pre_node_name = upsample_layer['name'] 418 419 elif block['type'] == 'cost': 420 continue 421 422 # spacetodepth 423 elif block['type'] == 'reorg': 424 input_shape = self.layer_map[pre_node_name].get_attr('_output_shape') 425 reorg_layer = OrderedDict() 426 reorg_layer['input'] = [pre_node_name] 427 if 'name' in block.keys(): 428 reorg_layer['name'] = block['name'] 429 else: 430 reorg_layer['name'] = 'layer%d-reorg' % i 431 432 reorg_layer['type'] = 'SpaceToDepth' 433 reorg_param = OrderedDict() 434 stride = int(block['stride']) 435 reorg_param['strides'] = stride 436 reorg_param['_output_shape'] = [-1, input_shape[1]/stride, input_shape[2]/stride, input_shape[3]*stride*stride] 437 reorg_layer['attr'] = reorg_param 438 439 self.layer_map[reorg_layer['name']] = DarknetGraphNode(reorg_layer) 440 self.original_list[reorg_layer['name']] = DarknetGraphNode(reorg_layer) 441 self.layer_num_map[i] = reorg_layer['name'] 442 pre_node_name = reorg_layer['name'] 443 444 445 elif block['type'] == 'region': 446 # print(block) 447 region_layer = OrderedDict() 448 region_layer['input'] = [pre_node_name] 449 if 'name' in block.keys(): 450 region_layer['name'] = block['name'] 451 else: 452 region_layer['name'] = 'layer%d-region' % i 453 region_layer['type'] = 'region' 454 region_param = OrderedDict() 455 region_param['softmax'] = int(block['softmax']) 456 region_param['thresh'] = float(block['thresh']) 457 region_param['random'] = float(block['random']) 458 region_param['jitter'] = float(block['jitter']) 459 region_param['num'] = int(block['num']) 460 region_param['classes'] = int(block['classes']) 461 region_param['coords'] = int(block['coords']) 462 region_param['rescore'] = int(block['rescore']) 463 region_param['object_scale'] = int(block['object_scale']) 464 465 region_param['noobject_scale'] = int(block['noobject_scale']) 466 region_param['class_scale'] = int(block['class_scale']) 467 region_param['coord_scale'] = int(block['coord_scale']) 468 469 region_param['bias_match'] = int(block['bias_match']) 470 region_param['absolute'] = int(block['absolute']) 471 472 anchors = [float(t) for t in block['anchors'].split(',')] 473 region_param['anchors'] = anchors 474 475 region_layer['attr'] = region_param 476 # print(region_layer) 477 self.layer_map[region_layer['name']] = DarknetGraphNode(region_layer) 478 self.original_list[region_layer['name']] = DarknetGraphNode(region_layer) 479 self.layer_num_map[i] = region_layer['name'] 480 # assert False 481 482 483 else: 484 print('unknown layer type %s ' % block['type']) 485 print(block,"\n") 486 assert False 487 488 489 490 for layer in self.layer_map: 491 for pred in self.layer_map[layer].layer['input']: 492 if pred not in self.layer_map.keys() and pred != 'data': 493 print(pred) 494 print("::::::::::::: unknown input :::::::::::::") 495 assert False 496 497 self._make_connection(pred, layer) 498 499 super(DarknetGraph, self).build() 500 501