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 transform_map_fn_output_concatenation(external_match: dict,
internal_match: dict):
"""
Transforms TensorFlow 2 output concatenation into use of axis attribute for output port of Loop node
:param external_match: a match used for handling a part of the main graph responsible for output concatenation
:param internal_match: a match used for handling a part of the body graph responsible for output concatenation
"""
loop_node = external_match['while']
stack_node = external_match['stack']
list_reserve_node = external_match['reserve']
body_graph = loop_node['body']
tensor_list_set_item_node = internal_match['concatenation']
tensor_list_set_item_node_name = tensor_list_set_item_node.soft_get(
'name', tensor_list_set_item_node.id)
list_result_node = internal_match['concatenation_result']
# replace TensorListSetItem with Unsqueeze and use axis attribute for corresponding Result node
# to concatenate results from different iterations
unsqueeze_list_element = create_op_with_const_inputs(
body_graph, Unsqueeze, {1: int64_array(0)},
{'name': 'TensorListSetItemUnsqueeze'})
tensor_list_set_item_node.in_port(2).get_connection().set_destination(
unsqueeze_list_element.in_port(0))
tensor_list_set_item_node.out_port(0).get_connection().set_source(
unsqueeze_list_element.out_port(0))
rename_nodes([(tensor_list_set_item_node,
tensor_list_set_item_node_name + '/AbandonedName'),
(unsqueeze_list_element, tensor_list_set_item_node_name)
])
list_result_node_layer_id = list_result_node.internal_layer_id
Loop.update_port_map_value_ext(loop_node.output_port_map,
'internal_layer_id',
list_result_node_layer_id, 'axis', 0)
# remove TensorListStack to by-pass the node since the result from the Loop node is already concatenated
stack_node.out_port(0).get_connection().set_source(
stack_node.in_port(0).get_connection().get_source())
# disconnect ListReserve node because it is no longer needed for Loop
list_reserve_node.out_port(0).disconnect()
# connect a number of iterations with trip count that can be received from the second input of ListReserve
# create a constant network with True value for execution_condition so that IE can ignore execution condition
# and perform trip_counts iterations. This approach with known trip count value allows to avoid dynamism.
loop_node.in_port(1).disconnect()
list_reserve_node.in_port(1).get_source().connect(loop_node.in_port(1))
for record in loop_node.output_port_map:
if 'purpose' in record and record[
'purpose'] == 'execution_condition':
exec_cond_layer_id = record['internal_layer_id']
exec_cond_node = Loop.get_body_node_by_internal_id(
loop_node, exec_cond_layer_id)
const_true = Const(body_graph, {
'value': np.array(True, dtype=np.bool)
}).create_node()
exec_cond_node.in_port(0).get_connection().set_source(
const_true.out_port(0))
