input = node.input[i] if input in id_dict: node.input[i] = id_dict[input] def generate_edges(self): for node in self.nodes: for input in node.input: if input in self.edges: self.edges[input].append(node) else: self.edges[input] = [node] @staticmethod def get_scope_name(name): index = name.rfind('/') if index == -1: return "" return name[0:index] def in_conv2d_scope(self, name): inner_scope = TFConverter.get_scope_name(name) if inner_scope == "": return False; for scope in self.conv2d_scope_names: index = inner_scope.find(scope) if index == 0: return True return False def in_dense_scope(self, name): inner_scope = TFConverter.get_scope_name(name) if inner_scope == "": return False; for scope in self.dense_scope_names: index = inner_scope.find(scope) if index == 0: return True return False def generate_sub_block_op_scope_info(self): # mostly, conv2d/dense is a sub block in graph, get the scope name for node in self.nodes: if node.op == 'Conv2D': scope = TFConverter.get_scope_name(node.name) # for the case tf.nn.conv2d is called directly if scope == '': continue # for the case tf.nn.conv2d is called within a scope if scope + '/kernel' not in self.name_node_dict: continue self.conv2d_scope_names.add(scope) elif node.op == 'MatMul': scope = TFConverter.get_scope_name(node.name) # for the case tf.nn.dense is called directly if scope == '': continue # for the case tf.nn.dense is called within a scope if scope + '/kernel' not in self.name_node_dict and scope.split('/Tensordot')[0] + '/kernel' not in self.name_node_dict: continue self.dense_scope_names.add(scope.split('/Tensordot')[0]) # get the input name to the conv2d/dense sub block for node in self.nodes: scope = TFConverter.get_scope_name(node.name) if scope in self.conv2d_scope_names: if node.op == 'Conv2D' or node.op == 'Shape': for inp in node.input: if TFConverter.get_scope_name(inp) != scope: self.conv2d_scopename_inputname_dict[scope] = inp elif scope in self.dense_scope_names: if node.op == 'MatMul' or node.op == 'Shape': for inp in node.input: if TFConverter.get_scope_name(inp) != scope: self.dense_scopename_inputname_dict[scope] = inp elif scope.split('/Tensordot')[0] in self.dense_scope_names: if node.op == 'Transpose': for inp in node.input: if TFConverter.get_scope_name(inp).find(scope)<0 and TFConverter.get_scope_name(inp).find(scope.split('/')[0])<0: self.dense_scopename_inputname_dict[scope.split('/Tensordot')[0]] = inp def run(self): self.generate_name_node_dict() self.generate_output_names() self.remove_identity() self.generate_edges() self.generate_sub_block_op_scope_info() if self.dump4tb: self.dump_for_tensorboard() self.dump_to_file() def convert_from_tensorflow(infile, outfile, dump4tb): with open(infile, 'rb') as f: # read the file in .proto format graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) nodes = graph_def.node converter = TFConverter(graph_def, nodes, outfile, dump4tb) converter.run()Lines 501–607 of 607. Previous lines
input = node.input[i] if input in id_dict: node.input[i] = id_dict[input] def generate_edges(self): for node in self.nodes: for input in node.input: if input in self.edges: self.edges[input].append(node) else: self.edges[input] = [node] @staticmethod def get_scope_name(name): index = name.rfind('/') if index == -1: return "" return name[0:index] def in_conv2d_scope(self, name): inner_scope = TFConverter.get_scope_name(name) if inner_scope == "": return False; for scope in self.conv2d_scope_names: index = inner_scope.find(scope) if index == 0: return True return False def in_dense_scope(self, name): inner_scope = TFConverter.get_scope_name(name) if inner_scope == "": return False; for scope in self.dense_scope_names: index = inner_scope.find(scope) if index == 0: return True return False def generate_sub_block_op_scope_info(self): # mostly, conv2d/dense is a sub block in graph, get the scope name for node in self.nodes: if node.op == 'Conv2D': scope = TFConverter.get_scope_name(node.name) # for the case tf.nn.conv2d is called directly if scope == '': continue # for the case tf.nn.conv2d is called within a scope if scope + '/kernel' not in self.name_node_dict: continue self.conv2d_scope_names.add(scope) elif node.op == 'MatMul': scope = TFConverter.get_scope_name(node.name) # for the case tf.nn.dense is called directly if scope == '': continue # for the case tf.nn.dense is called within a scope if scope + '/kernel' not in self.name_node_dict and scope.split('/Tensordot')[0] + '/kernel' not in self.name_node_dict: continue self.dense_scope_names.add(scope.split('/Tensordot')[0]) # get the input name to the conv2d/dense sub block for node in self.nodes: scope = TFConverter.get_scope_name(node.name) if scope in self.conv2d_scope_names: if node.op == 'Conv2D' or node.op == 'Shape': for inp in node.input: if TFConverter.get_scope_name(inp) != scope: self.conv2d_scopename_inputname_dict[scope] = inp elif scope in self.dense_scope_names: if node.op == 'MatMul' or node.op == 'Shape': for inp in node.input: if TFConverter.get_scope_name(inp) != scope: self.dense_scopename_inputname_dict[scope] = inp elif scope.split('/Tensordot')[0] in self.dense_scope_names: if node.op == 'Transpose': for inp in node.input: if TFConverter.get_scope_name(inp).find(scope)<0 and TFConverter.get_scope_name(inp).find(scope.split('/')[0])<0: self.dense_scopename_inputname_dict[scope.split('/Tensordot')[0]] = inp def run(self): self.generate_name_node_dict() self.generate_output_names() self.remove_identity() self.generate_edges() self.generate_sub_block_op_scope_info() if self.dump4tb: self.dump_for_tensorboard() self.dump_to_file() def convert_from_tensorflow(infile, outfile, dump4tb): with open(infile, 'rb') as f: # read the file in .proto format graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) nodes = graph_def.node converter = TFConverter(graph_def, nodes, outfile, dump4tb) converter.run()