Showing the changes to tools/python (1 of 9 files). Show all
Unified · Split
| @@ -48,9 +48,9 @@ class Operand(object): · 47 unmodified lines | ||
| 48 | 48 | self.used_count = self.used_count + 1 |
| 49 | 49 | |
| 50 | 50 | def __str__(self): |
| 51 | − return "{}: (name: {}, iotype: {}, dtype: {}, dims: ({},{},{},{}) used_count: {})".format(self.index, | |
| 51 | + return "{}: (name: {}, iotype: {}, dtype: {}, dims: {}, used_count: {})".format(self.index, | |
| 52 | 52 | self.name, self.iotype2str[self.iotype], self.dtype2str[self.dtype], |
| 53 | − self.dims[0], self.dims[1], self.dims[2], self.dims[3], self.used_count) | |
| 53 | + self.dims, self.used_count) | |
| 54 | 54 | |
| 55 | 55 | def __lt__(self, other): |
| 56 | 56 | return self.index < other.index |
| @@ -71,8 +71,10 @@ class TFConverter: · 14 unmodified lines | ||
| 71 | 71 | self.converted_nodes = set() |
| 72 | 72 | self.conv2d_scope_names = set() |
| 73 | 73 | self.conv2d_scopename_inputname_dict = {} |
| 74 | + self.dense_scope_names = set() | |
| 75 | + self.dense_scopename_inputname_dict = {} | |
| 74 | 76 | self.op2code = {'Conv2D':1, 'DepthToSpace':2, 'MirrorPad':3, 'Maximum':4, |
| 75 | − 'MathBinary':5, 'MathUnary':6, 'AvgPool':7} | |
| 77 | + 'MathBinary':5, 'MathUnary':6, 'AvgPool':7, 'MatMul':8} | |
| 76 | 78 | self.mathbin2code = {'Sub':0, 'Add':1, 'Mul':2, 'RealDiv':3, 'Minimum':4, 'FloorMod':5} |
| 77 | 79 | self.mathun2code = {'Abs':0, 'Sin':1, 'Cos':2, 'Tan':3, 'Asin':4, |
| 78 | 80 | 'Acos':5, 'Atan':6, 'Sinh':7, 'Cosh':8, 'Tanh':9, 'Asinh':10, |
| @@ -126,6 +128,22 @@ class TFConverter: · 47 unmodified lines | ||
| 126 | 128 | return knode, bnode, dnode, anode |
| 127 | 129 | |
| 128 | 130 | |
| 131 | + def get_dense_params(self, dense_scope_name): | |
| 132 | + knode = self.name_node_dict[dense_scope_name + '/kernel'] | |
| 133 | + bnode = self.name_node_dict.get(dense_scope_name + '/bias') | |
| 134 | + # the BiasAdd name is possible be changed into the output name, | |
| 135 | + # if activation is None, and BiasAdd.next is the last op which is Identity | |
| 136 | + anode = None | |
| 137 | + if bnode: | |
| 138 | + if dense_scope_name + '/BiasAdd' in self.edges: | |
| 139 | + anode = self.edges[dense_scope_name + '/BiasAdd'][0] | |
| 140 | + if anode.op not in self.conv_activations: | |
| 141 | + anode = None | |
| 142 | + else: | |
| 143 | + anode = None | |
| 144 | + return knode, bnode, anode | |
| 145 | + | |
| 146 | + | |
| 129 | 147 | def dump_complex_conv2d_to_file(self, node, f): |
| 130 | 148 | assert(node.op == 'Conv2D') |
| 131 | 149 | self.layer_number = self.layer_number + 1 |
| @@ -181,6 +199,57 @@ class TFConverter: · 49 unmodified lines | ||
| 181 | 199 | output_operand_index = self.add_operand(self.edges[bnode.name][0].name, Operand.IOTYPE_OUTPUT) |
| 182 | 200 | np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f) |
| 183 | 201 | |
| 202 | + def dump_dense_to_file(self, node, f): | |
| 203 | + assert(node.op == 'MatMul') | |
| 204 | + self.layer_number = self.layer_number + 1 | |
| 205 | + self.converted_nodes.add(node.name) | |
| 206 | + | |
| 207 | + scope_name = TFConverter.get_scope_name(node.name) | |
| 208 | + #knode for kernel, bnode for bias, anode for activation | |
| 209 | + knode, bnode, anode = self.get_dense_params(scope_name.split('/')[0]) | |
| 210 | + | |
| 211 | + if bnode is not None: | |
| 212 | + has_bias = 1 | |
| 213 | + btensor = bnode.attr['value'].tensor | |
| 214 | + if btensor.tensor_shape.dim[0].size == 1: | |
| 215 | + bias = struct.pack("f", btensor.float_val[0]) | |
| 216 | + else: | |
| 217 | + bias = btensor.tensor_content | |
| 218 | + else: | |
| 219 | + has_bias = 0 | |
| 220 | + | |
| 221 | + if anode is not None: | |
| 222 | + activation = anode.op | |
| 223 | + else: | |
| 224 | + activation = 'None' | |
| 225 | + | |
| 226 | + ktensor = knode.attr['value'].tensor | |
| 227 | + in_channels = ktensor.tensor_shape.dim[0].size | |
| 228 | + out_channels = ktensor.tensor_shape.dim[1].size | |
| 229 | + if in_channels * out_channels == 1: | |
| 230 | + kernel = np.float32(ktensor.float_val[0]) | |
| 231 | + else: | |
| 232 | + kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32) | |
| 233 | + kernel = kernel.reshape(in_channels, out_channels) | |
| 234 | + kernel = np.transpose(kernel, [1, 0]) | |
| 235 | + | |
| 236 | + np.array([self.op2code[node.op], self.conv_activations[activation], in_channels, out_channels, has_bias], dtype=np.uint32).tofile(f) | |
| 237 | + kernel.tofile(f) | |
| 238 | + if has_bias: | |
| 239 | + f.write(bias) | |
| 240 | + | |
| 241 | + input_name = self.dense_scopename_inputname_dict[scope_name.split('/')[0]] | |
| 242 | + input_operand_index = self.add_operand(input_name, Operand.IOTYPE_INPUT) | |
| 243 | + | |
| 244 | + if anode is not None: | |
| 245 | + output_operand_index = self.add_operand(anode.name, Operand.IOTYPE_OUTPUT) | |
| 246 | + else: | |
| 247 | + if bnode is not None: | |
| 248 | + output_operand_index = self.add_operand(self.edges[bnode.name][0].name, Operand.IOTYPE_OUTPUT) | |
| 249 | + else: | |
| 250 | + output_operand_index = self.add_operand(self.edges[scope_name+'/concat_1'][0].name, Operand.IOTYPE_OUTPUT) | |
| 251 | + np.array([input_operand_index, output_operand_index], dtype=np.uint32).tofile(f) | |
| 252 | + | |
| 184 | 253 | |
| 185 | 254 | def dump_simple_conv2d_to_file(self, node, f): |
| 186 | 255 | assert(node.op == 'Conv2D') |
| @@ -343,9 +412,19 @@ class TFConverter: · 156 unmodified lines | ||
| 343 | 412 | if node.op == 'Conv2D': |
| 344 | 413 | self.dump_complex_conv2d_to_file(node, f) |
| 345 | 414 | continue |
| 415 | + if self.in_dense_scope(node.name): | |
| 416 | + if node.op == 'MatMul': | |
| 417 | + self.dump_dense_to_file(node, f) | |
| 418 | + continue | |
| 419 | + | |
| 346 | 420 | |
| 347 | 421 | if node.op == 'Conv2D': |
| 348 | 422 | self.dump_simple_conv2d_to_file(node, f) |
| 423 | + continue | |
| 424 | + if node.name in self.output_names: | |
| 425 | + input_name = self.id_different_scope_dict[node.name] | |
| 426 | + if TFConverter.get_scope_name(input_name)!=TFConverter.get_scope_name(node.name): | |
| 427 | + continue | |
| 349 | 428 | if node.op == 'AvgPool': |
| 350 | 429 | self.dump_avg_pool_to_file(node, f) |
| 351 | 430 | elif node.op == 'DepthToSpace': |
| @@ -367,7 +446,7 @@ class TFConverter: · 15 unmodified lines | ||
| 367 | 446 | np.array([operand.index, len(operand.name)], dtype=np.uint32).tofile(f) |
| 368 | 447 | f.write(operand.name.encode('utf-8')) |
| 369 | 448 | np.array([operand.iotype, operand.dtype], dtype=np.uint32).tofile(f) |
| 370 | − np.array([operand.dims[0], operand.dims[1], operand.dims[2], operand.dims[3]], dtype=np.uint32).tofile(f) | |
| 449 | + np.array(operand.dims, dtype=np.uint32).tofile(f) | |
| 371 | 450 | |
| 372 | 451 | |
| 373 | 452 | def dump_to_file(self): |
| @@ -396,6 +475,7 @@ class TFConverter: · 22 unmodified lines | ||
| 396 | 475 | |
| 397 | 476 | |
| 398 | 477 | def remove_identity(self): |
| 478 | + self.id_different_scope_dict = {} | |
| 399 | 479 | id_nodes = [] |
| 400 | 480 | id_dict = {} |
| 401 | 481 | for node in self.nodes: |
| @@ -408,6 +488,7 @@ class TFConverter: · 6 unmodified lines | ||
| 408 | 488 | self.name_node_dict[input].name = name |
| 409 | 489 | self.name_node_dict[name] = self.name_node_dict[input] |
| 410 | 490 | del self.name_node_dict[input] |
| 491 | + self.id_different_scope_dict[name] = input | |
| 411 | 492 | else: |
| 412 | 493 | id_dict[name] = input |
| 413 | 494 | |
| @@ -449,8 +530,18 @@ class TFConverter: · 35 unmodified lines | ||
| 449 | 530 | return False |
| 450 | 531 | |
| 451 | 532 | |
| 452 | − def generate_conv2d_scope_info(self): | |
| 453 | − # mostly, conv2d is a sub block in graph, get the scope name | |
| 533 | + def in_dense_scope(self, name): | |
| 534 | + inner_scope = TFConverter.get_scope_name(name) | |
| 535 | + if inner_scope == "": | |
| 536 | + return False; | |
| 537 | + for scope in self.dense_scope_names: | |
| 538 | + index = inner_scope.find(scope) | |
| 539 | + if index == 0: | |
| 540 | + return True | |
| 541 | + return False | |
| 542 | + | |
| 543 | + def generate_sub_block_op_scope_info(self): | |
| 544 | + # mostly, conv2d/dense is a sub block in graph, get the scope name | |
| 454 | 545 | for node in self.nodes: |
| 455 | 546 | if node.op == 'Conv2D': |
| 456 | 547 | scope = TFConverter.get_scope_name(node.name) |
| @@ -461,8 +552,17 @@ class TFConverter: · 4 unmodified lines | ||
| 461 | 552 | if scope + '/kernel' not in self.name_node_dict: |
| 462 | 553 | continue |
| 463 | 554 | self.conv2d_scope_names.add(scope) |
| 555 | + elif node.op == 'MatMul': | |
| 556 | + scope = TFConverter.get_scope_name(node.name) | |
| 557 | + # for the case tf.nn.dense is called directly | |
| 558 | + if scope == '': | |
| 559 | + continue | |
| 560 | + # for the case tf.nn.dense is called within a scope | |
| 561 | + if scope + '/kernel' not in self.name_node_dict and scope.split('/Tensordot')[0] + '/kernel' not in self.name_node_dict: | |
| 562 | + continue | |
| 563 | + self.dense_scope_names.add(scope.split('/Tensordot')[0]) | |
| 464 | 564 | |
| 465 | − # get the input name to the conv2d sub block | |
| 565 | + # get the input name to the conv2d/dense sub block | |
| 466 | 566 | for node in self.nodes: |
| 467 | 567 | scope = TFConverter.get_scope_name(node.name) |
| 468 | 568 | if scope in self.conv2d_scope_names: |
| @@ -470,6 +570,16 @@ class TFConverter: · 1 unmodified line | ||
| 470 | 570 | for inp in node.input: |
| 471 | 571 | if TFConverter.get_scope_name(inp) != scope: |
| 472 | 572 | self.conv2d_scopename_inputname_dict[scope] = inp |
| 573 | + elif scope in self.dense_scope_names: | |
| 574 | + if node.op == 'MatMul' or node.op == 'Shape': | |
| 575 | + for inp in node.input: | |
| 576 | + if TFConverter.get_scope_name(inp) != scope: | |
| 577 | + self.dense_scopename_inputname_dict[scope] = inp | |
| 578 | + elif scope.split('/Tensordot')[0] in self.dense_scope_names: | |
| 579 | + if node.op == 'Transpose': | |
| 580 | + for inp in node.input: | |
| 581 | + if TFConverter.get_scope_name(inp).find(scope)<0 and TFConverter.get_scope_name(inp).find(scope.split('/')[0])<0: | |
| 582 | + self.dense_scopename_inputname_dict[scope.split('/Tensordot')[0]] = inp | |
| 473 | 583 | |
| 474 | 584 | |
| 475 | 585 | def run(self): |
| @@ -477,7 +587,7 @@ class TFConverter: · 1 unmodified line | ||
| 477 | 587 | self.generate_output_names() |
| 478 | 588 | self.remove_identity() |
| 479 | 589 | self.generate_edges() |
| 480 | − self.generate_conv2d_scope_info() | |
| 590 | + self.generate_sub_block_op_scope_info() | |
| 481 | 591 | |
| 482 | 592 | if self.dump4tb: |
| 483 | 593 | self.dump_for_tensorboard() |