def extract(cls, loop_node): Loop.update_node_stat(loop_node, {}) # check that required body and condition functions exist in the graph library main_graph = loop_node.graph body_graph_proto = get_graph_proto(main_graph, 'body', loop_node) cond_graph_proto = get_graph_proto(main_graph, 'cond', loop_node) body_graph = create_internal_graph(main_graph) loop_node['body'] = body_graph # create Parameter nodes for the body graph body_parameters, body_parameter_names = convert_graph_inputs_to_parameters( body_graph, body_graph_proto) # 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 = {} for attr_key, attr_value in main_graph.graph.items(): if attr_key not in ['tensor_mapping', 'parent_node']: main_graph_attrs_copy[attr_key] = copy.deepcopy(attr_value) body_graph.graph.update(main_graph_attrs_copy) loop_node['body'] = body_graph # save parent node for nested loops to know which node contains body (and which graph is on upper level) body_graph.graph['parent_node'] = loop_node # 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 # save additional edges information for graph on each level, the first one is the deepest 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') if hasattr(body_graph, 'op_names_statistic') and hasattr(pb_node, 'op_type'): body_graph.op_names_statistic[pb_node.op_type] += 1 # add incoming edges based on data_nodes_map for dst_port, inp in enumerate(pb_node.input): # should add edge src_internal_id --> dst_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 else: is_finished = create_cross_body_edge(body_graph, external_edges, additional_params, inp, id, dst_port) if not is_finished: 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] create_edge_with_attrs(body_graph, inp, src_id, src_port, id, dst_port) # 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 # 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 cur_graph = body_graph for external_edges_subg in external_edges: if 'parent_node' not in cur_graph.graph: continue cur_loop_node = cur_graph.graph['parent_node'] parent_graph = cur_loop_node.graph for (src_node, src_port), body_node, tensor_name in external_edges_subg: create_edge_with_attrs(parent_graph, tensor_name, src_node, src_port, cur_loop_node.id, next_loop_input_port_idx) Loop.connect_body_input(cur_loop_node, next_loop_input_port_idx, body_node) next_loop_input_port_idx += 1 cur_graph = parent_graph # 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 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): Loop.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): Loop.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): Loop.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