def normalize_body_graph(loop_node: Node):
loop_name = loop_node.soft_get('name', loop_node.id)
# connect "trip count" input if it is not connected with default value "Infinity" (-1)
if not loop_node.is_in_port_connected(0):
loop_node.add_input_port(0, skip_if_exist=True)
Const(loop_node.graph, {'name': loop_name + '/trip_count', 'value': int64_array(-1)}).\
create_node().out_port(0).connect(loop_node.in_port(0))
# connect "execution condition" input if it is not connected with default value True
if not loop_node.is_in_port_connected(1):
loop_node.add_input_port(1, skip_if_exist=True)
Const(loop_node.graph, {'name': loop_name + '/execution_cond', 'value': np.array(True, dtype=np.bool)}).\
create_node().out_port(0).connect(loop_node.in_port(1))
# scan output need Unsqueeze over axis 0
for record in loop_node.output_port_map:
body_node = Loop.get_body_node_by_internal_id(loop_node, record['internal_layer_id'])
assert body_node is not None
assert body_node.soft_get('type') == 'Result'
if record['axis'] is not None:
unsqueeze = create_op_with_const_inputs(loop_node.body, Unsqueeze, {1: int64_array([0])})
body_node.in_port(0).get_connection().insert_node(unsqueeze)
Loop.normalize_input_output_ports(loop_node)
def normalize_loop_node(graph: Graph, loop_node: Node):
loop_name = loop_node.soft_get('name', loop_node.id)
# disconnect current iteration from external port #0 and move trip count to this port
loop_node.in_port(0).disconnect()
loop_node.in_port(1).get_connection().add_destination(loop_node.in_port(0))
Loop.update_port_map_value(loop_node.input_port_map, 'external_port_id', 1, 0)
# connect execution condition port
exec_cond_node = Const(graph, {'name': loop_name + '/ExecutionConditionValue',
'value': np.array(True, dtype=np.bool)}).create_node()
loop_node.in_port(1).get_connection().set_source(exec_cond_node.out_port(0))
loop_node.body.clean_up()
Loop.normalize_input_output_ports(loop_node)
def find_and_replace_pattern(self, graph: Graph):
cleanup_called_once = False
# walk through all Loop nodes and find Const inputs
for loop_node in graph.get_op_nodes(op='Loop'):
# call clean-up only once that performs constant folding
if not cleanup_called_once:
graph.clean_up()
cleanup_called_once = True
# move constant node into the body graph and removes body parameter nodes corresponding to them
Loop.pull_constant_inputs_into_body(loop_node)
# since some input ports can be removed after the pulling constants, normalization of Loop node is required
Loop.normalize_input_output_ports(loop_node)
# perform shape inference for the Loop node again since new constant can be appeared
# and constant folding can be helpful for weights path to Convolution node inside the body graph
loop_node['need_shape_inference'] = True
