def extract(cls, node): attrs = { 'shape': node.shape, 'data_type': np.float32, # TODO: other types? } Parameter.update_node_stat(node, {}) return cls.enabled
def extract(node): attrs = { 'data_type': tf_dtype_extractor(node.pb.attr["dtype"].type), 'shape': tf_tensor_shape(node.pb.attr["shape"].shape), 'permute_attrs': PermuteAttrs().update_attrs(attrs=[('shape', 'output:0')]) } Parameter.update_node_stat(node, attrs) return __class__.enabled
def extract(node): attrs = { 'shape': np.array([d.dim_value for d in node.pb.type.tensor_type.shape.dim], dtype=np.int64), } Parameter.update_node_stat(node, attrs) return __class__.enabled
def create_parameter_with_empty_attrs(graph, param_name): graph.add_node(param_name, kind='op', op='Parameter', name=param_name, pb=None, shape=None) parameter_node = Node(graph, param_name) # need to manually update necessary attrs for the node because extractor will not be called # for it because the node does not have .pb attribute Parameter.update_node_stat(parameter_node, {}) parameter_node['internal_layer_id'] = len(graph.nodes) return parameter_node
def extract(cls, node): t_type = node.pb.type.tensor_type attrs = { 'shape': np.array([d.dim_value for d in t_type.shape.dim], dtype=np.int64), 'data_type': TENSOR_TYPE_TO_NP_TYPE[t_type.elem_type] } Parameter.update_node_stat(node, attrs) return cls.enabled
def extract(cls, node): t_type = node.pb.type.tensor_type attrs = { 'shape': shape_array([ d.dim_value if (not hasattr(d, 'dim_param') or d.dim_param == '') and d.dim_value != 0 else dynamic_dimension_value for d in t_type.shape.dim ]), 'data_type': TENSOR_TYPE_TO_NP_TYPE[t_type.elem_type] } Parameter.update_node_stat(node, attrs) return cls.enabled
def convert_graph_inputs_to_parameters(internal_graph, internal_graph_proto): # create Parameter nodes for the body graph body_parameters = [] body_parameter_names = [] for idx, pb_node in enumerate(internal_graph_proto['input_arg']): param_id = internal_graph.unique_id(pb_node.name) internal_graph.add_node(param_id, name=param_id, kind='op', op='Parameter', pb=None, shape=None) parameter_node = Node(internal_graph, pb_node.name) Parameter.update_node_stat( parameter_node, { 'data_type': tf_dtype_extractor(pb_node.type), 'permute_attrs': PermuteAttrs().update_attrs(attrs=[('shape', 'output:0')]) }) body_parameters.append(parameter_node) body_parameter_names.append(param_id) return body_parameters, body_parameter_names
def extract(cls, node): if 'value' in node.symbol_dict: Const.update_node_stat(node, {'value': node.symbol_dict['value']}) else: Parameter.update_node_stat(node, {}) return cls.enabled
def extract(cls, loop_node): Loop.update_node_stat(loop_node, {}) loop_name = loop_node.soft_get('name', loop_node.id) # check that required body and condition functions exist in the graph library main_graph = loop_node.graph body_graph_name = loop_node.pb.attr['body'].func.name cond_graph_name = loop_node.pb.attr['cond'].func.name assert 'library' in main_graph.graph, 'The graph does not contain a library that is required ' \ 'by node with name "{}".'.format(loop_name) library_graph = main_graph.graph['library'] assert body_graph_name in library_graph, 'The library does not contain a function with name "{}" ' \ 'that is required by node ' \ 'with name "{}".'.format(body_graph_name, loop_name) body_graph_proto = library_graph[body_graph_name] assert cond_graph_name in library_graph, 'The library does not contain a function with name "{}" ' \ 'that is required by node ' \ 'with name "{}".'.format(cond_graph_name, loop_name) cond_graph_proto = library_graph[cond_graph_name] body_graph = Graph() # fill the body graph for attr_key in main_graph.graph.keys(): if attr_key != 'library': body_graph.graph[attr_key] = copy.deepcopy(main_graph.graph[attr_key]) else: # it is sufficient to have a link to the library body_graph.graph['library'] = main_graph.graph['library'] loop_node['body'] = body_graph # create Parameter nodes for the body graph body_parameters = [] body_parameter_names = [] for idx, pb_node in enumerate(body_graph_proto['input_arg']): param_id = body_graph.unique_id(pb_node.name) body_graph.add_node(param_id, name=param_id, kind='op', op='Parameter', pb=None, shape=None) parameter_node = Node(body_graph, pb_node.name) Parameter.update_node_stat(parameter_node, {'data_type': tf_dtype_extractor(pb_node.type), 'permute_attrs': PermuteAttrs().update_attrs(attrs=[('shape', 'output:0')])} ) body_parameters.append(parameter_node) body_parameter_names.append(param_id) # update the loop body graph with the body function graph body_results = [] update_body_graph(body_graph, body_graph_proto, body_parameter_names, body_results) # update the loop body graph with the condition function graph update_body_graph(body_graph, cond_graph_proto, body_parameter_names, body_results) # add 'internal_layer_id' attribute which is a must have attribute for the loop body node for idx, body_node in enumerate(body_graph.get_op_nodes()): body_node['internal_layer_id'] = idx body_graph.stage = 'front' # Currently, # Loop Inputs Order: # 0 - current iteration # 1 - trip count # 2.. - "loop carried" dependencies variables # # Body Inputs Order: # 0 - current iteration # 1 - trip count # 2.. - "loop carried" dependencies variables # # Body Outputs Order: # 0 - current iteration # 1 - trip count # 2.. - "loop carried" dependencies variables # # Loop Outputs Order: # 0 - current iteration # 1 - trip count # 2.. - "loop carried" dependencies variables # # so inputs must be reordered and execution condition must be created in the front transformation # to be aligned with the specification # connect external input ports with body parameter nodes except current iteration # since it must be disconnected from external port for idx in range(1, len(body_parameters)): Loop.connect_body_input(loop_node, idx, body_parameters[idx]) # mark current iteration input Parameter node and execution condition Result node Loop.mark_current_iteration_parameter_node(loop_node, body_parameters[0]) Loop.mark_execution_condition_result_node(loop_node, body_results[-1]) # connect back edges in the body except current iteration for idx in range(1, len(body_parameters)): Loop.add_back_edge(loop_node, body_parameters[idx], body_results[idx]) # connect body outputs with Loop operation output ports except the execution condition result for idx in range(len(body_results)-1): Loop.connect_body_output(loop_node, idx, body_results[idx]) # run function to parse body nodes attributes similar to the main graph extract_node_attrs(body_graph, lambda node: tf_op_extractor(node, check_for_duplicates(tf_op_extractors))) return cls.enabled
def extract(cls, loop_node): Loop.update_node_stat(loop_node, {}) body_graph_proto = onnx_attr(loop_node, 'body', 'g', None) main_graph = loop_node.graph # create a Graph object for the body and take graph attributes from the main graph body_graph = Graph() main_graph_attrs_copy = copy.deepcopy(main_graph.graph) del main_graph_attrs_copy['tensor_mapping'] body_graph.graph.update(main_graph_attrs_copy) loop_node['body'] = body_graph # maps a tensor name to a node produced it and the node port: str -> (node_id, node_port) data_nodes_map = {} body_graph.graph[ 'tensor_mapping'] = data_nodes_map # save mapping for possible Loop inside the Loop body_parameters = add_initializers_and_inputs_to_graph( body_graph, body_graph_proto, data_nodes_map) external_edges = [ ] # (src_node, src_out_port), dest_body_parameter_node additional_params = { } # (src_node, src_out_port) -> parameter_node (for manually added Parameters) # Go through all nodes in the original model order because data nodes are defined on-the-fly and order matters for pb_node in body_graph_proto.node: # create an NX node id = body_graph.unique_id(node_id(pb_node)) body_graph.add_node(id, pb=pb_node, kind='op') # add incoming edges based on data_nodes_map for dst_port, inp in enumerate(pb_node.input): # should add edge inp --> id if inp not in data_nodes_map: if inp == '': # input is omitted; most likely it corresponds to an optional input for an operator continue elif inp in main_graph.graph['tensor_mapping']: log.debug( 'The edge between outer and inner graphs detected: {} -> {}' .format(inp, id)) if main_graph.graph['tensor_mapping'][ inp] not in additional_params: # create new Parameter body node and connect the body node with the outer graph using it param_id = str(inp) body_graph.add_node(param_id, kind='op', op='Parameter', name=param_id, pb=None, shape=None) parameter_node = Node(body_graph, param_id) # need to manually update necessary attrs for the node because extractor will not be called # for it because the node does not have .pb attribute Parameter.update_node_stat(parameter_node, {}) external_edges.append( (main_graph.graph['tensor_mapping'][inp], parameter_node)) src_id, src_port = param_id, 0 additional_params[main_graph.graph[ 'tensor_mapping'][inp]] = parameter_node else: src_id, src_port = additional_params[ main_graph.graph['tensor_mapping'][inp]].id, 0 else: raise Error( 'Reference to "{}" is not satisfied. A node refer not existing data tensor. ONNX ' 'model is not consistent. Protobuf fragment: {}', inp, pb_node) else: src_id, src_port = data_nodes_map[inp] assert (body_graph.has_node(src_id)) edge_attrs = { 'out': src_port, 'in': dst_port, 'name': inp, 'fw_tensor_debug_info': [(inp, inp)], 'in_attrs': ['in', 'name'], 'out_attrs': ['out', 'name'], 'data_attrs': ['fw_tensor_debug_info'] } body_graph.add_edge(src_id, id, **edge_attrs) # add outgoing edges to data_nodes_map for src_port, out in enumerate(pb_node.output): if out in data_nodes_map: log.debug("Detected reuse of blob {}.".format(out)) data_nodes_map[out] = (id, src_port) body_results = [] for output in body_graph_proto.output: tensor_name = str(output.name) node_name, output_port = data_nodes_map[tensor_name] assert body_graph.has_node( node_name ), 'The body graph does not contain output with name "{}"'.format( node_name) body_results.append( Node(body_graph, add_opoutput(body_graph, node_name, output_port, False))) # add 'internal_layer_id' attribute which is a must have attribute for the loop body node for idx, body_node in enumerate(body_graph.get_op_nodes()): body_node['internal_layer_id'] = idx loop_carried_dependencies_count = len(body_graph_proto.input) - 2 scan_outputs_count = len( body_graph_proto.output) - 1 - loop_carried_dependencies_count # Loop inputs: # 0 - trip count # 1 - execution condition # 2 .. - loop carried dependencies # Loop outputs: # 0 .. loop_carried_dependencies_count - 1 - loop carried dependencies # loop_carried_dependencies_count .. - scan outputs # Body inputs: # 0 - iteration number # 1 - execution condition # 2 .. - loop carried dependencies # Body outputs: # 0 - execution condition # 1 .. loop_carried_dependencies_count - loop carried dependencies # loop_carried_dependencies_count + 1 .. - scan outputs body_graph.stage = 'front' # some of the inputs/outputs may not be connected but the normalization transformation will take care of it # connection Loop body nodes with external input edges next_loop_input_port_idx = sorted(loop_node.in_edges().keys())[-1] + 1 for (src_node, src_port), body_node in external_edges: main_graph.add_edge( src_node, loop_node.id, **{ 'out': src_port, 'in': next_loop_input_port_idx, 'name': src_node, 'fw_tensor_debug_info': [(src_node, src_node)], 'in_attrs': ['in', 'name'], 'out_attrs': ['out', 'name'], 'data_attrs': ['fw_tensor_debug_info'] }) connect_body_input(loop_node, next_loop_input_port_idx, body_node) next_loop_input_port_idx += 1 # mark current iteration input Parameter node Loop.mark_current_iteration_parameter_node(loop_node, body_parameters[0]) # connect initial value for "execution condition" input of the loop connect_body_input(loop_node, 1, body_parameters[1]) # add back edge with "execution condition" Loop.add_back_edge(loop_node, body_parameters[1], body_results[0]) # mark "execution condition" Result node Loop.mark_execution_condition_result_node(loop_node, body_results[0]) # connect initial value for "loop carried" dependencies variables for idx in range(loop_carried_dependencies_count): connect_body_input(loop_node, idx + 2, body_parameters[idx + 2]) # add back edge for "loop carried" dependencies variables for idx in range(loop_carried_dependencies_count): Loop.add_back_edge(loop_node, body_parameters[idx + 2], body_results[idx + 1]) # connect final value for "loop carried" dependencies variables for idx in range(loop_carried_dependencies_count): connect_body_output(loop_node, idx, body_results[idx + 1]) # connect "scan outputs" and mark axis for concatenation for idx in range(loop_carried_dependencies_count, loop_carried_dependencies_count + scan_outputs_count): connect_body_output(loop_node, idx, body_results[idx + 1], axis=0) # run function to parse body nodes attributes similar to the main graph extract_node_attrs( body_graph, lambda node: onnx_op_extractor( node, check_for_duplicates(onnx_op_extractors))) return cls.enabled
def extract(node): Parameter.update_node_stat(node) return __class__.enabled
def extract(node): Parameter.update_node_stat( node, {'shape': dim_to_shape(node.pb.input_param.shape[0].dim)}) return __class__.enabled
def extract(cls, node): Parameter.update_node_stat(node) return cls.enabled