def replace_pattern(graph: nx.MultiDiGraph, match: dict): reshape_node = match['reshape'] in_node = reshape_node.in_node() out_node = reshape_node.out_node() if not np.array_equal(in_node.shape, out_node.shape): return False remove_op_node_with_data_node(graph, reshape_node)
def replace_pattern(self, graph: Graph, match: dict): ti = match['ti'] direct_reverse = match['direct_reverse'] inverse_reverse = match['inverse_reverse'] assert direct_reverse.seq_axis == inverse_reverse.seq_axis assert direct_reverse.batch_axis is None and inverse_reverse.batch_axis is None or \ direct_reverse.batch_axis == inverse_reverse.batch_axis # Modify stride in TI for port_map in [ti.input_port_map, ti.output_port_map]: for port in port_map: if 'axis' in port and port[ 'axis'] is not None and 'external_port_id' in port: assert port['axis'] == direct_reverse.seq_axis, \ 'axis == {} != {} == direct_reverse.seq_dim'.format(port['axis'], direct_reverse.seq_axis) if 'stride' not in port or port['stride'] is None: port['stride'] = 1 assert port['stride'] in [-1, 1] port['stride'] = -port['stride'] if port['stride'] == -1: port['start'] = -1 port['end'] = 0 elif port['stride'] == 1: port['start'] = None port['end'] = None # Remove reverses remove_op_node_with_data_node(graph, direct_reverse) remove_op_node_with_data_node(graph, inverse_reverse)
def find_and_replace_pattern(self, graph: Graph): for node in graph.pseudo_topological_sort(): if node.kind == 'data' or node.op != 'Switch': continue switch_op_node = node pred_id_data_node = switch_op_node.in_node(1) graph.remove_edge(pred_id_data_node.id, switch_op_node.id) remove_op_node_with_data_node(graph, switch_op_node)
def replace_pattern(graph: Graph, match: dict): """ Removes output SoftMax layer :param graph: graph to operate on :param match: dictionary with matched nodes """ if len(match['softmax_data'].out_nodes()) == 1: remove_op_node_with_data_node(graph, match['softmax_node'])
def find_and_replace_pattern(self, graph: Graph): for n in pseudo_topological_sort(graph): if graph.node[n]['kind'] == 'data' or graph.node[n]['op'] != 'Switch': continue switch_op_node = Node(graph, n) pred_id_data_node = switch_op_node.in_node(1) graph.remove_edge(pred_id_data_node.id, switch_op_node.id) remove_op_node_with_data_node(graph, switch_op_node)
def replace_pattern(self, graph: Graph, match: dict): if len(graph.in_edges(match['merge'].id)) <= 1: remove_op_node_with_data_node( graph, match['merge'], list(match['merge'].in_nodes().values())[0]) log.info( "Useles Merge op and data nodes was deleted op='{}'".format( match['merge'].id))
def replace_pattern(graph: nx.MultiDiGraph, match: dict): node = match['reshape_1'] if (node.has_valid('type') and node.type == 'Reshape' and len(node.out_nodes()) == 1 and node.out_node().has_valid('kind') and node.out_node().kind == 'data' and len(node.out_node().out_nodes()) == 1): remove_op_node_with_data_node(graph, node)
def replace_pattern(self, graph: Graph, match: dict): node = match['pad'] for port, input_node in node.in_nodes().items(): if port != 0: graph.remove_edge(input_node.id, node.id) # remove Pad operation if all pads are equal to 0 if np.all(node.pads == 0): remove_op_node_with_data_node(graph, node)
def replace_pattern(graph: Graph, match: dict): """ Removes output SoftMax layer :param graph: graph to operate on :param match: dictionary with matched nodes """ if len(match['softmax_data'].out_nodes()) == 1: remove_op_node_with_data_node(graph, match['softmax_node']) else: log.error("SoftMax is not last layer, so can't be removed", extra={'is_warning': True})
def find_and_replace_pattern(self, graph: Graph): for node in graph.get_op_nodes(op='Pad'): all_pads_zeros = True for in_port_ind in range(1, 3): input_node = node.in_port(in_port_ind).get_source().node value = input_node.soft_get('value', None) all_pads_zeros &= input_node.soft_get('type') == 'Const' and value is not None and np.all(value == 0) if all_pads_zeros: remove_op_node_with_data_node(graph, node)
def replace_pattern(graph: Graph, match: dict): node_ss = match['strided_slice'] # slices = [elem for elem in node_ss.slices if elem is not None] # node_ss.slices = np.array(slices) if node_ss.out_port(0).data.get_value() is not None: # StridedSlices(SS) in shape-calculating sub-graphs that should not be deleted that easily # Example: # In RetinaNetFilteredDetectionsReplacement we have SS that slices first batch # We delete such SS for batch 1, but it should be performed while reshaping the model return output_data_node = node_ss.out_node(0) input_data_node = node_ss.in_node(0) out_shape = output_data_node.shape if not np.all(node_ss.shrink_axis_mask == 0): out_shape = list(out_shape) for i in range(len(node_ss.shrink_axis_mask)): if node_ss.shrink_axis_mask[i] == 1: out_shape.insert(i, 1) out_shape = int64_array(out_shape) if not np.all(node_ss.new_axis_mask == 0): out_shape = list(out_shape) for i in reversed(range(len(node_ss.new_axis_mask))): if node_ss.new_axis_mask[i] == 1: out_shape.pop(i) out_shape = int64_array(out_shape) if np.array_equal(input_data_node.shape, out_shape) and \ all(elem.step == 1 for elem in match['strided_slice'].slices): if not np.all(node_ss.shrink_axis_mask == 0): ConvertGroupedStridedSlice.add_squeeze_for_shrink( graph, node_ss) if not np.all(node_ss.new_axis_mask == 0): ConvertGroupedStridedSlice.add_unsqueeze_for_new( graph, node_ss) log.info("Useless StridedSlice op '{}' has been detected".format( match['strided_slice'].id)) # remove inputs to Strided Slice so it has just one input with data so we can use 'remove_op_node' function graph.remove_edge(match['strided_slice'].in_node(1).id, match['strided_slice'].id) graph.remove_edge(match['strided_slice'].in_node(2).id, match['strided_slice'].id) if len(match['strided_slice'].in_nodes()) > 3: graph.remove_edge(match['strided_slice'].in_node(3).id, match['strided_slice'].id) remove_op_node_with_data_node(graph, match['strided_slice'])
def fuse_sequence_of_reshapes(graph: Graph): for node in list(graph.nodes()): if not graph.has_node(node): # data node can be already removed continue node = Node(graph, node) if (node.has_valid('type') and node.type == 'Reshape' and len(node.out_nodes()) == 1 and node.out_node().has_valid('kind') and node.out_node().kind == 'data' and len(node.out_node().out_nodes()) == 1): log.debug('First phase for Reshape: {}'.format(node.name)) next_op = node.out_node().out_node() log.debug('second node: {}'.format(next_op.graph.node[next_op.id])) if next_op.has_valid('type') and next_op.type == 'Reshape': # Detected Reshape1 --> data --> Reshape2 pattern without side edges # Remove Reshape1 log.debug('Second phase for Reshape: {}'.format(node.name)) remove_op_node_with_data_node(graph, node) reshape_nodes = graph.get_op_nodes(op='Reshape') for reshape_node in reshape_nodes: in_ports = [ port for port in reshape_node.in_ports().values() if not port.disconnected() ] assert len(in_ports) in [ 1, 2 ], "`Reshape` node must have 2 inputs or 1 input with `dim`" if len(in_ports) == 2: previous_dim_op = reshape_node.in_port(1).get_source().node.op if previous_dim_op != 'Const': continue dim = reshape_node.in_port(1).get_connection().data.get_value() else: assert reshape_node.has_valid( 'dim'), "`Reshape` node with 1 input must have `dim` attribute" dim = reshape_node.dim in_shape = reshape_node.in_port(0).get_connection().data.get_shape() if np.array_equal(dim, in_shape) and len(reshape_node.out_nodes()): log.debug("Useless reshape with dim {} was deleted: {}".format( str(dim), reshape_node.name)) reshape_node.out_port(0).get_connection().set_source( reshape_node.in_port(0).get_source())
def replace_pattern(self, graph: nx.MultiDiGraph, match: dict): output_data_node = match['strided_slice'].out_node(0) input_data_node = match['strided_slice'].in_node(0) if np.array_equal(input_data_node.shape, output_data_node.shape) and \ all(elem.step == 1 for elem in match['strided_slice'].slices): log.info("Useless StridedSlice op '{}' has been detected".format( match['strided_slice'].id)) # remove inputs to Strided Slice so it has just one input with data so we can use 'remove_op_node' function graph.remove_edge(match['strided_slice'].in_node(1).id, match['strided_slice'].id) graph.remove_edge(match['strided_slice'].in_node(2).id, match['strided_slice'].id) graph.remove_edge(match['strided_slice'].in_node(3).id, match['strided_slice'].id) remove_op_node_with_data_node(graph, match['strided_slice'])
def replace_pattern(graph: nx.MultiDiGraph, match: dict): """ Need to find the pattern: Parent (any type) -> SoftMAx -> OpOutput It is needed to remove output SoftMAx layer Parameters ---------- graph : nx.MultiDiGraph Graph with loaded model. match : dict Patterns which were found in graph structure. """ softmax = match['softmax_node'] child = softmax.out_node() if not child.has_and_set('is_output'): return remove_op_node_with_data_node(graph, softmax)
def replace_pattern(graph: Graph, match: dict): node_ss = match['strided_slice'] output_data_node = node_ss.out_node(0) input_data_node = node_ss.in_node(0) out_shape = output_data_node.shape if not np.all(node_ss.shrink_axis_mask == 0): out_shape = list(out_shape) for i in range(len(node_ss.shrink_axis_mask)): if node_ss.shrink_axis_mask[i] == 1: out_shape.insert(i, 1) out_shape = int64_array(out_shape) if not np.all(node_ss.new_axis_mask == 0): out_shape = list(out_shape) for i in reversed(range(len(node_ss.new_axis_mask))): if node_ss.new_axis_mask[i] == 1: out_shape.pop(i) out_shape = int64_array(out_shape) if np.array_equal(input_data_node.shape, out_shape) and \ all(elem.step == 1 for elem in match['strided_slice'].slices): if not np.all(node_ss.shrink_axis_mask == 0): ConvertGroupedStridedSlice.add_squeeze_for_shrink( graph, node_ss) if not np.all(node_ss.new_axis_mask == 0): ConvertGroupedStridedSlice.add_unsqueeze_for_new( graph, node_ss) log.info("Useless StridedSlice op '{}' has been detected".format( match['strided_slice'].id)) # remove inputs to Strided Slice so it has just one input with data so we can use 'remove_op_node' function graph.remove_edge(match['strided_slice'].in_node(1).id, match['strided_slice'].id) graph.remove_edge(match['strided_slice'].in_node(2).id, match['strided_slice'].id) if len(match['strided_slice'].in_nodes()) > 3: graph.remove_edge(match['strided_slice'].in_node(3).id, match['strided_slice'].id) remove_op_node_with_data_node(graph, match['strided_slice'])
def fuse_sequence_of_reshapes(graph: nx.MultiDiGraph): for node in list(graph.nodes()): node = Node(graph, node) if not graph.has_node(node.id): # data node can be already removed continue if (node.has_valid('type') and node.type == 'Reshape' and len(node.out_nodes()) == 1 and node.out_node().has_valid('kind') and node.out_node().kind == 'data' and len(node.out_node().out_nodes()) == 1): log.debug('First phase for Reshape: {}'.format(node.name)) next_op = node.out_node().out_node() log.debug('second node: {}'.format(next_op.graph.node[next_op.id])) if next_op.has_valid('type') and next_op.type == 'Reshape': # Detected Reshape1 --> data --> Reshape2 pattern without side edges # Remove Reshape1 log.debug('Second phase for Reshape: {}'.format(node.name)) remove_op_node_with_data_node(graph, node)
def find_and_replace_pattern(self, graph: Graph): intervals = {} for node in graph.get_op_nodes(type='FakeQuantize', keep_in_IR=False): prev_node = node.in_node().in_node() prev_node_id = prev_node.id prev_node_out_shape = prev_node.out_node()['shape'] C = prev_node_out_shape[1] assert node.in_node(1).value.size == 1 assert node.in_node(2).value.size == 1 # Input and output ranges should match if we want to remove FakeQuantize from model assert_msg = "FakeQuantize cannot be removed because input and output intervals do not match" assert node.in_node(1).value == node.in_node(3).value, assert_msg assert node.in_node(2).value == node.in_node(4).value, assert_msg min = ', '.join([str(node.in_node(1).value.flatten()[0])] * C) max = ', '.join([str(node.in_node(2).value.flatten()[0])] * C) intervals[prev_node_id] = {'min': min, 'max': max} remove_op_node_with_data_node(graph, node) if intervals: if 'statistics' not in graph.graph: graph.graph['statistics'] = intervals else: graph.graph['statistics'].update(intervals)
def replace_pattern(self, graph: Graph, match: dict): ti = match['ti'] direct_reverse = match['direct_reverse'] inverse_reverse = match['inverse_reverse'] assert direct_reverse.seq_axis == inverse_reverse.seq_axis assert direct_reverse.batch_axis is None and inverse_reverse.batch_axis is None or \ direct_reverse.batch_axis == inverse_reverse.batch_axis if not self.is_fusable_reverse_sequence(direct_reverse) or \ not self.is_fusable_reverse_sequence(inverse_reverse): # we can not merge ReverseSequence without equal sequences return # Modify stride in TI for port_map in [ti.input_port_map, ti.output_port_map]: for port in port_map: if 'axis' in port and port[ 'axis'] is not None and 'external_port_id' in port: assert port['axis'] == direct_reverse.seq_axis, \ 'axis == {} != {} == direct_reverse.seq_dim'.format(port['axis'], direct_reverse.seq_axis) if 'stride' not in port or port['stride'] is None: port['stride'] = 1 assert port['stride'] in [-1, 1] port['stride'] = -port['stride'] if port['stride'] == -1: port['start'] = -1 port['end'] = 0 elif port['stride'] == 1: port['start'] = 0 port['end'] = -1 # disconnect subgraph for seq length calculation direct_reverse.in_port(1).disconnect() inverse_reverse.in_port(1).disconnect() # Remove reverses remove_op_node_with_data_node(graph, direct_reverse) remove_op_node_with_data_node(graph, inverse_reverse)
def find_and_replace_pattern(self, graph: Graph): reshape_nodes = graph.get_op_nodes(type='Reshape') for node in reshape_nodes: if not graph.has_node(node.id): # the Reshape node has been removed in the previous iteration continue if len(node.out_port(0).get_destinations()) == 1: log.debug('First phase for Reshape: {}'.format( node.soft_get('name'))) next_op = get_next_operation(node)[0] log.debug('second node: id={}, type={}'.format( next_op.soft_get('id'), next_op.soft_get('type'))) if next_op.has_valid('type') and next_op.type == 'Reshape': dim_value = next_op.in_port(1).data.get_value() if dim_value is None or 0 in dim_value or -1 in dim_value: # we do not fuse reshape sequences with special symbols: 0, -1 continue # Detected Reshape1 --> data --> Reshape2 pattern without side edges. Remove Reshape1 log.debug('Second phase for Reshape: {}'.format( node.soft_get('name'))) remove_op_node_with_data_node(graph, node)
def replace_pattern(self, graph: Graph, match: dict): quantize = match['quantize'] # Check for total number of ReLU consumers -- if something else consume its output it cannot be fused if len(match['relu'].out_node().out_nodes()) > 1: log.debug('ReluQuantizeFuse: cannot fuse because ReLU have multiple consumers') return # If the fusion is applicable, direct modifications to quantize 1-st and 2-nd inputs # are performed. So the data nodes at those inputs shouldn't have more than 1 consumer # maximum 2 consumers to the same quantize op (consumed by 1st and 2nd ports). # TODO: relax this limitation and duplicate data nodes accordingly to modify the input range freely # Provisional limitation that related to binary quantization # TODO: Relax it beyond binarization case if len(quantize.in_node(1).out_nodes()) != 2 or \ len(quantize.in_node(2).out_nodes()) != 2 or \ quantize.in_node(1).id != quantize.in_node(2).id or \ quantize.levels != 2: log.debug('ReluQuantizeFuse: cannot fuse because Quantize op has ' 'unexpected number of consumers for ports 1 and 2') return threshold = quantize.in_node(1) # As we restricted to binarization case only, so we need to detect from # which side of 0 Quantize threshold resides: # if the threshold > 0, it remains the same; # if the threshold == 0, it also remains the same; # if the threshold < 0, it should be modified to -infinity that means that all inputs map to output_high modification_mask = threshold.value < 0 threshold.value[modification_mask] = float('-inf') # Remove ReLU as it no longer needed remove_op_node_with_data_node(graph, match['relu'])
def replace_pattern(self, graph: Graph, match: dict): # This transformation works if and only if a body of TI # matches the following topology (Reshape -> LSTMCell -> Reshape) nodes = [('input_unsqueezed'), ('squeeze', dict(op='Reshape')), ('input_squeezed'), ('input_hidden'), ('input_cell'), ('weights'), ('biases'), ('lstm', dict(op='LSTMCell')), ('output_hidden'), ('output_cell'), ('unsqueeze', dict(op='Reshape')), ('output_unsqueezed'), ('const_w', dict(op='Const')), ('const_b', dict(op='Const')), ('op_output', dict(op='OpOutput')), ('op_output_1', dict(op='OpOutput')), ('op_output_2', dict(op='OpOutput'))] edges = [ ('input_unsqueezed', 'squeeze'), ('squeeze', 'input_squeezed'), ('input_squeezed', 'lstm', { 'in': 0 }), ('input_hidden', 'lstm', { 'in': 1 }), ('input_cell', 'lstm', { 'in': 2 }), ('weights', 'lstm', { 'in': 3 }), ('biases', 'lstm', { 'in': 4 }), ('const_w', 'weights'), ('const_b', 'biases'), ('lstm', 'output_hidden', { 'out': 0 }), ('lstm', 'output_cell', { 'out': 1 }), ('output_hidden', 'unsqueeze'), ('unsqueeze', 'output_unsqueezed'), ('output_unsqueezed', 'op_output'), ('output_hidden', 'op_output_1'), ('output_cell', 'op_output_2'), ] ti = match['ti'] isomorphisms = find_isomorphisms(ti.body, nodes, edges) if len(list(isomorphisms)) != 1: return isomorphism = isomorphisms[0] direct_permute = match['direct_permute'] inverse_permute = match['inverse_permute'] permute_order = [1, 0, 2] # Check both perumute orders exactly match expected one - [1, 0, 2] if not direct_permute.has_valid('order') or not np.array_equal( direct_permute.order, permute_order): return if not inverse_permute.has_valid('order') or not np.array_equal( inverse_permute.order, permute_order): return def find_ports(port_map: list, attrs: dict): """ Find all ports in a given port map with specified attributes """ result = [] for i, port in enumerate(port_map): if dict_includes(port, attrs): result.append(i) return result # Check TI has only single partitioned input/output port; all partitioned ports have defined axis data_input_port = find_ports(ti.input_port_map, {'axis': lambda attr: attr in [0, 1]}) data_output_port = find_ports(ti.output_port_map, {'axis': lambda attr: attr in [0, 1]}) assert len(data_input_port) == 1 assert len(data_output_port) == 1 data_input_port = data_input_port[0] data_output_port = data_output_port[0] # Verify that they are really connected to Permute layers (guarantied by port numbers of TI, see the pattern) assert ti.in_edge(0)['external_port_id'] == ti.input_port_map[ data_input_port]['external_port_id'] assert ti.out_edge(0)['external_port_id'] == ti.output_port_map[ data_output_port]['external_port_id'] # Verify that the TI body have required Reshapes connected to the found ports squeeze = isomorphism['squeeze'] unsqueeze = isomorphism['unsqueeze'] assert squeeze['internal_layer_id'] == ti.input_port_map[ data_input_port]['internal_layer_id'] assert squeeze.in_edge(0)['internal_port_id'] == ti.input_port_map[ data_input_port]['internal_port_id'] assert unsqueeze['internal_layer_id'] == ti.output_port_map[ data_output_port]['internal_layer_id'] assert unsqueeze.out_edge(0)['internal_port_id'] == ti.output_port_map[ data_output_port]['internal_port_id'] assert len(squeeze.in_node().shape) == 3 assert len(squeeze.out_node().shape) == 2 assert len(unsqueeze.in_node().shape) == 2 assert len(unsqueeze.out_node().shape) == 3 # Remove permutes remove_op_node_with_data_node(graph, direct_permute) remove_op_node_with_data_node(graph, inverse_permute) match['output'].shape = match['output'].shape[permute_order] # swap 0/1 axis for partitioned ports ti.input_port_map[data_input_port][ 'axis'] = 1 - ti.input_port_map[data_input_port]['axis'] ti.output_port_map[data_output_port][ 'axis'] = 1 - ti.output_port_map[data_output_port]['axis'] # smap 0-th and 1-th shape entries for reshapes inside body squeeze.in_node().shape = squeeze.in_node().shape[[1, 0, 2]] unsqueeze.out_node().shape = unsqueeze.out_node().shape[[1, 0, 2]] unsqueeze.dim = unsqueeze.dim[[1, 0, 2]]
def replace_pattern(self, graph: Graph, match: dict): remove_op_node_with_data_node(graph, match['op'])
def replace_pattern(self, graph: Graph, match: dict): # This transformation works if and only if a body of TI # matches the following topology (Squeeze -> LSTMCell -> Unsqueeze) nodes = [ ('squeeze_dim', dict(kind='op', op='Const')), ('squeeze_dim_data', dict(kind='data')), ('unsqueeze_dim', dict(kind='op', op='Const')), ('unsqueeze_dim_data', dict(kind='data')), ('input_unsqueezed', dict(kind='data')), ('squeeze', dict(kind='op', op='Squeeze')), ('input_squeezed', dict(kind='data')), ('input_hidden', dict(kind='data')), ('input_cell', dict(kind='data')), ('weights', dict(kind='data')), ('biases', dict(kind='data')), ('lstm', dict(kind='op', op='LSTMCell')), ('output_hidden', dict(kind='data')), ('output_cell', dict(kind='data')), ('unsqueeze', dict(kind='op', op='Unsqueeze')), ('output_unsqueezed', dict(kind='data')), ('const_w', dict(kind='op', op='Const')), ('const_b', dict(kind='op', op='Const')), ('op_output', dict(kind='op', op='Result')), ('op_output_1', dict(kind='op', op='Result')), ('op_output_2', dict(kind='op', op='Result')), ('input_unsqueezed_i', dict(kind='op', op='Parameter')), ('input_hidden_i', dict(kind='op', op='Parameter')), ('input_cell_i', dict(kind='op', op='Parameter')), ] edges = [ ('input_unsqueezed', 'squeeze', { 'in': 0 }), ('squeeze', 'input_squeezed'), ('squeeze_dim', 'squeeze_dim_data'), ('squeeze_dim_data', 'squeeze', { 'in': 1 }), ('input_squeezed', 'lstm', { 'in': 0 }), ('input_hidden', 'lstm', { 'in': 1 }), ('input_cell', 'lstm', { 'in': 2 }), ('weights', 'lstm', { 'in': 3 }), ('biases', 'lstm', { 'in': 4 }), ('const_w', 'weights'), ('const_b', 'biases'), ('lstm', 'output_hidden', { 'out': 0 }), ('lstm', 'output_cell', { 'out': 1 }), ('output_hidden', 'unsqueeze'), ('unsqueeze', 'output_unsqueezed'), ('unsqueeze_dim', 'unsqueeze_dim_data'), ('unsqueeze_dim_data', 'unsqueeze', { 'in': 1 }), ('output_unsqueezed', 'op_output'), ('output_hidden', 'op_output_1'), ('output_cell', 'op_output_2'), ('input_unsqueezed_i', 'input_unsqueezed'), ('input_hidden_i', 'input_hidden'), ('input_cell_i', 'input_cell'), ] ti = match['ti'] isomorphisms = find_isomorphisms(ti.body, nodes, edges) if len(list(isomorphisms)) != 1: return isomorphism = isomorphisms[0] direct_permute = match['direct_permute'] inverse_permute = match['inverse_permute'] permute_order = [1, 0, 2] # Check both perumute orders exactly match expected one - [1, 0, 2] direct_order = direct_permute.in_port(1).data.get_value() if direct_order is None or not np.array_equal(direct_order, permute_order): return inverse_order = inverse_permute.in_port(1).data.get_value() if inverse_order is None or not np.array_equal(inverse_order, permute_order): return # Check non-ShapeOf output out of direct Transpose is exactly one direct_permute_dsts = direct_permute.out_port(0).get_destinations() if len([ dst for dst in direct_permute_dsts if dst.node.soft_get('type') != 'ShapeOf' ]) != 1: return for shape_of_dst in [ dst for dst in direct_permute_dsts if dst.node.soft_get('type') == 'ShapeOf' ]: name = shape_of_dst.node.soft_get( 'name', shape_of_dst.node.id) + '/FusedToTITranspose' gather = create_op_with_const_inputs( graph, op=Gather, op_attrs={'name': name}, port_value_dict={ 1: int64_array(permute_order), 2: int64_array(0) }) shape_of_dst.node.out_port(0).get_connection().insert_node(gather) def find_ports(port_map: list, attrs: dict): """ Find all ports in a given port map with specified attributes """ result = [] for i, port in enumerate(port_map): if dict_includes(port, attrs): result.append(i) return result # Check TI has only single partitioned input/output port; all partitioned ports have defined axis data_input_port = find_ports(ti.input_port_map, {'axis': lambda attr: attr in [0, 1]}) data_output_port = find_ports(ti.output_port_map, {'axis': lambda attr: attr in [0, 1]}) assert len(data_input_port) == 1 assert len(data_output_port) == 1 data_input_port = data_input_port[0] data_output_port = data_output_port[0] # Verify that they are really connected to Transpose layers (guarantied by port numbers of TI, see the pattern) assert ti.in_edge(0)['external_port_id'] == ti.input_port_map[ data_input_port]['external_port_id'] assert ti.out_edge(0)['external_port_id'] == ti.output_port_map[ data_output_port]['external_port_id'] # Verify that the TI body have required Reshapes connected to the found ports squeeze = isomorphism['squeeze'] unsqueeze = isomorphism['unsqueeze'] assert len(squeeze.in_node().shape) == 3 assert len(squeeze.out_node().shape) == 2 assert len(unsqueeze.in_node().shape) == 2 assert len(unsqueeze.out_node().shape) == 3 # Remove permutes remove_op_node_with_data_node(graph, direct_permute) remove_op_node_with_data_node(graph, inverse_permute) match['output'].shape = match['output'].shape[permute_order] # swap 0/1 axis for partitioned ports ti.input_port_map[data_input_port][ 'axis'] = 1 - ti.input_port_map[data_input_port]['axis'] ti.output_port_map[data_output_port][ 'axis'] = 1 - ti.output_port_map[data_output_port]['axis'] isomorphism['input_unsqueezed_i'].shape = isomorphism[ 'input_unsqueezed_i'].shape[[1, 0, 2]] isomorphism['input_unsqueezed_i'].infer( isomorphism['input_unsqueezed_i']) isomorphism['squeeze_dim'].value = ti.input_port_map[data_input_port][ 'axis'] isomorphism['squeeze_dim'].infer(isomorphism['squeeze_dim']) isomorphism['squeeze']['need_shape_inference'] = True isomorphism['unsqueeze_dim'].value = ti.output_port_map[ data_output_port]['axis'] isomorphism['unsqueeze_dim'].infer(isomorphism['unsqueeze_dim']) isomorphism['unsqueeze'].infer(isomorphism['unsqueeze'])