import copy import functools import itertools import re import typing import warnings import igraph as ig import numpy as np from tinynn.util.util import class_conditional, get_logger from ..schemas.tflite.schema_generated import ActivationFunctionType, Padding from . import tflite as tfl from .base import FUSE_ACTIVATION_MAP, ExtendedOperator from .graph import CommonGraph log = get_logger(__name__, 'INFO') class GraphOptimizer(object): graph: CommonGraph fuse_tensor_count: int fuse_attr_count: int fuse_quant: bool group_conv_rewrite: bool tflite_micro_rewrite: bool quantize_input_output_type: typing.Optional[str] # Optimization levels NO_OPTIMIZE: int = 0 FOLD_BUFFER: int = 1 FUSE_BN: int = 2 COMMON_OPTIMIZE: int = 3 BRANCH_OPTIMIZE: int = 4 BRANCH_OPTIMIZE_EXTENDED: int = 5 ALL_OPTIMIZE: int = 5 def __init__( self, graph: CommonGraph, level: int, fuse_quant: bool, group_conv_rewrite: bool, rewrite_quantizable: bool, tflite_micro_rewrite: bool, quantize_input_output_type: typing.Optional[str], fuse_input_indices: typing.Optional[typing.List[int]] = None, fuse_output_indices: typing.Optional[typing.List[int]] = None, max_transpose_dims: int = -1, bypass_elementwise_passthrough_constraint: bool = False, group_tensors: bool = False, conv_transpose_with_bias: bool = True, hybrid_int16_lstm: bool = False, ) -> None: self.graph = graph self.fuse_tensor_count = 0 self.fuse_attr_count = 0 self.level = level self.fuse_quant = fuse_quant self.group_conv_rewrite = group_conv_rewrite self.rewrite_quantizable = rewrite_quantizable self.tflite_micro_rewrite = tflite_micro_rewrite self.quantize_input_output_type = quantize_input_output_type self.fuse_input_indices = fuse_input_indices self.fuse_output_indices = fuse_output_indices self.max_transpose_dims = max_transpose_dims self.bypass_elementwise_passthrough_constraint = bypass_elementwise_passthrough_constraint self.group_tensors = group_tensors self.conv_transpose_with_bias = conv_transpose_with_bias self.hybrid_int16_lstm = hybrid_int16_lstm def create_attr_tensor( self, tensor: tfl.Tensor, name: str = None, quantization: typing.Optional[tfl.QuantizationParameters] = None ): if name is None: if self.fuse_attr_count == 0: name = 'fuse_attr' else: name = f'fuse_attr_{self.fuse_attr_count}' self.fuse_attr_count += 1 return tfl.Tensor(tensor, name, has_buffer=True, quantization=quantization) def create_transform_tensor( self, tensor: tfl.Tensor, name: str = None, quantization: typing.Optional[tfl.QuantizationParameters] = None ): if name is None: if self.fuse_tensor_count == 0: name = 'fuse_transform' else: name = f'fuse_transform_{self.fuse_tensor_count}' self.fuse_tensor_count += 1 return tfl.Tensor(tensor, name, has_buffer=False, quantization=quantization) @class_conditional(lambda self: self.level >= GraphOptimizer.FUSE_BN) def fuse_conv_fc_bn(self): # Find fusable ops edges = self.graph.graph.es.select(functools.partial(is_bn_fusable_edge, graph_converter=self.graph.graph)) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target], x) for x in edges) remove_ids = [] actions = [] for conv, bn, tensor in filtered_pairs: bn_activ = bn['op'].fusedActivationFunction conv_activ = getattr(conv['op'], 'fusedActivationFunction', None) if conv_activ is None and bn_activ != ActivationFunctionType.NONE: continue # Find out the output of the batch-norm nodes new_output = bn['outputs'][0] assert new_output in self.graph.tensor_map # For each node that is next of a batch-norm node, we connect it with the conv node self.graph.connect_next_tensors(bn, conv, new_output) # Update graph, prepare to drop the output tensor of the conv node and use the output tensor of the # batch-norm instead conv['outputs'][0] = new_output conv['op'].outputs[0] = self.graph.tensor_map[new_output] self.graph.tensor_node_map[new_output] = conv['name'] tensor['name'] = bn['outputs'][0] tensor['label'] = bn['outputs'][0] if bn_activ != ActivationFunctionType.NONE and conv_activ == ActivationFunctionType.NONE: conv['op'].fusedActivationFunction = bn_activ # Collect the arguments of the conv and batch-norm nodes weight = conv['op'].inputs[1] bias = conv['op'].inputs[2] if len(conv['op'].inputs) > 2 else None bn_w, bn_b, bn_mean, bn_var = bn['op'].inputs[1:] bn_w, bn_b, bn_mean, bn_var = ( bn_w.tensor.copy(), bn_b.tensor.copy(), bn_mean.tensor.copy(), bn_var.tensor.copy(), ) activ_w = weight.tensor.copy() activ_b = bias.tensor.copy() if bias is not None else None eps = bn['op'].eps # Fuse conv/fc and batch-norm new_weight = fuse_bn_weight( eps, bn_w, bn_var, activ_w, conv['node_type'] == ExtendedOperator.GENERIC_DECONV ) new_bias = fuse_bn_bias(eps, bn_w, bn_var, bn_mean, bn_b, activ_b) # New attribute tensors new_w = self.create_attr_tensor(new_weight) new_b = self.create_attr_tensor(new_bias) # Collect the actions we should take here # The reason that we don't do the actions here is because we are currently in the loop of vertices, # the iterator will be invalidated once `replace_operator_input` is called actions.append((self.graph.replace_operator_input, (conv, 1, new_w))) if bias is not None: actions.append((self.graph.replace_operator_input, (conv, 2, new_b))) else: actions.append((self.graph.append_operator_input, (conv, new_b))) remove_ids.append(bn.index) # Process actions for func, args in actions: func(*args) # Delete batch-norm nodes for id in remove_ids: vertex = self.graph.graph.vs[id] assert vertex['node_type'] == ExtendedOperator.BATCH_NORM self.graph.graph.delete_vertices(remove_ids) @class_conditional(lambda self: self.level >= GraphOptimizer.FUSE_BN) def fuse_bn_conv(self): edges = self.graph.graph.es.select(functools.partial(is_rev_bn_fusable_edge, graph_converter=self.graph.graph)) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]) for x in edges) def _remove_last_pred(seq): bn = seq[0] conv = seq[1] # Collect the arguments of the conv and batch-norm nodes weight = conv['op'].inputs[1] bias = conv['op'].inputs[2] if len(conv['op'].inputs) > 2 else None bn_w, bn_b, bn_mean, bn_var = bn['op'].inputs[1:] bn_w, bn_b, bn_mean, bn_var = ( bn_w.tensor.copy(), bn_b.tensor.copy(), bn_mean.tensor.copy(), bn_var.tensor.copy(), ) activ_w = weight.tensor.copy() activ_b = bias.tensor.copy() if bias is not None else None eps = bn['op'].eps new_weight = fuse_rev_bn_weight(eps, bn_w, bn_var, activ_w) new_bias = fuse_rev_bn_bias(eps, bn_w, bn_var, bn_mean, bn_b, activ_b, activ_w) return False, (conv, bias, new_weight, new_bias) def _remove_last_action(first_node, last_node, custom_data): conv, bias, new_weight, new_bias = custom_data new_w = self.create_attr_tensor(new_weight) new_b = self.create_attr_tensor(new_bias) actions = [] actions.append((self.graph.replace_operator_input, (conv, 1, new_w))) if bias is not None: actions.append((self.graph.replace_operator_input, (conv, 2, new_b))) else: actions.append((self.graph.append_operator_input, (conv, new_b))) return actions def _skip_pred(seq): bn = seq[0]['op'] conv = seq[1]['op'] skip = bn.inputs[0].quantization is not None or ( conv.inputs[1].shape[1] == 1 and conv.inputs[1].shape[0] == conv.groups and conv.groups > 1 ) return skip elinimate_sequences( self.graph, filtered_pairs, True, None, _remove_last_pred, _remove_last_action, _skip_pred, force_forward_input=True, ) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_activation(self): # Find fusable ops edges = self.graph.graph.es.select(functools.partial(is_activ_fusable_edge, graph_converter=self.graph.graph)) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target], x) for x in edges) remove_ids = [] for pre_activ, activ, tensor in filtered_pairs: if not self.conv_transpose_with_bias and pre_activ['node_type'] == ExtendedOperator.GENERIC_DECONV: continue # Find out the output of the batch-norm nodes new_output = activ['outputs'][0] assert new_output in self.graph.tensor_map # For each node that is next of the activation node, we connect it with the previous node self.graph.connect_next_tensors(activ, pre_activ, new_output) # Update graph, prepare to drop the output tensor of the conv node and use the output tensor of the # batch-norm instead pre_activ['outputs'][0] = new_output pre_activ['op'].outputs[0] = self.graph.tensor_map[new_output] self.graph.tensor_node_map[new_output] = pre_activ['name'] tensor['name'] = activ['outputs'][0] tensor['label'] = activ['outputs'][0] # Fuse activation pre_activ['op'].fusedActivationFunction = FUSE_ACTIVATION_MAP[activ['node_type']] remove_ids.append(activ.index) # Delete activation nodes self.graph.graph.delete_vertices(remove_ids) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_same_padding(self): edges = self.graph.graph.es.select(functools.partial(is_padding_fusable_edge, graph_converter=self.graph.graph)) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]) for x in edges) def _remove_last_pred(seq): op = seq[1]['op'] return False, op def _remove_last_action(first_node, last_node, custom_data): op = custom_data op.padding = Padding.SAME return [] def _skip_pred(seq): pad_op = seq[0]['op'] next_op = seq[1]['op'] input_shape = pad_op.inputs[0].shape[1:-1] if seq[1]['node_type'] == ExtendedOperator.MAX_POOL_2D: kernel_shape = (next_op.filterHeight, next_op.filterWidth) strides = (next_op.strideH, next_op.strideW) dilation = (1, 1) elif seq[1]['node_type'] in ( ExtendedOperator.CONV_2D, ExtendedOperator.DEPTHWISE_CONV_2D, ): kernel_shape = next_op.inputs[1].shape[1:-1] strides = (next_op.strideH, next_op.strideW) dilation = (next_op.dilationHFactor, next_op.dilationWFactor) elif seq[1]['node_type'] == ExtendedOperator.CONV_3D: kernel_shape = next_op.inputs[1].shape[:3] strides = (next_op.strideD, next_op.strideH, next_op.strideW) dilation = (next_op.dilationDFactor, next_op.dilationHFactor, next_op.dilationWFactor) pad_args = get_same_padding_args(input_shape, kernel_shape, strides, dilation) pad_arr = np.array(pad_args, dtype='int32') old_pad_arr = pad_op.inputs[1].tensor skip = not np.array_equal(pad_arr, old_pad_arr) return skip elinimate_sequences( self.graph, filtered_pairs, True, None, _remove_last_pred, _remove_last_action, _skip_pred, force_forward_input=True, ) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_same_padding_slicing(self): edges = self.graph.graph.es.select(functools.partial(is_slicing_fusable_edge, graph_converter=self.graph.graph)) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target], x) for x in edges) remove_ids = [] actions = [] for prev_node, slice_node, tensor in filtered_pairs: prev_op = prev_node['op'] slice_op = slice_node['op'] input_shape = slice_op.outputs[0].shape[1:-1] if prev_node['node_type'] == ExtendedOperator.TRANSPOSE_CONV: kernel_shape = prev_op.inputs[1].shape[1:-1] strides = (prev_op.strideH, prev_op.strideW) dilation = (1, 1) elif prev_node['node_type'] == ExtendedOperator.CONV_3D_TRANSPOSE: kernel_shape = prev_op.inputs[1].shape[:3] strides = (prev_op.strideD, prev_op.strideH, prev_op.strideW) dilation = (prev_op.dilationDFactor, prev_op.dilationHFactor, prev_op.dilationWFactor) pad_args = get_same_padding_args(input_shape, kernel_shape, strides, dilation) pad_arr = np.array(pad_args, dtype='int32') start_arr = [x for x in slice_op.inputs[1].tensor] end_arr = [slice_op.inputs[0].shape[i] - x - slice_op.outputs[0].shape[i] for i, x in enumerate(start_arr)] old_pad_args = [[x, y] for x, y in zip(start_arr, end_arr)] skip = not np.array_equal(pad_arr, old_pad_args) if skip: continue # Find out the output of the slice nodes new_output = slice_node['outputs'][0] assert new_output in self.graph.tensor_map # For each node that is next of the slice_nodeation node, we connect it with the previous node self.graph.connect_next_tensors(slice_node, prev_node, new_output) # Update graph, prepare to drop the output tensor of the conv node and use the output tensor of the # slice op instead prev_node['outputs'][0] = new_output prev_node['op'].outputs[0] = self.graph.tensor_map[new_output] self.graph.tensor_node_map[new_output] = prev_node['name'] tensor['name'] = slice_node['outputs'][0] tensor['label'] = slice_node['outputs'][0] # Fuse padding prev_node['op'].padding = Padding.SAME new_shape = np.array(prev_node['op'].outputs[0].shape, dtype='int32') new_shape_tensor = self.create_attr_tensor(new_shape) actions.append((self.graph.replace_operator_input, (prev_node, 0, new_shape_tensor))) remove_ids.append(slice_node.index) for func, args in actions: func(*args) # Delete activation nodes self.graph.graph.delete_vertices(remove_ids) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_requantize(self): # Find fusable ops edges = self.graph.graph.es.select( functools.partial(is_requantize_fusable_edge, graph_converter=self.graph.graph) ) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target], x) for x in edges) remove_ids = [] for pre_activ, activ, tensor in filtered_pairs: if pre_activ.outdegree() > 1: skip = False pre_quantize = None for out_edge in pre_activ.out_edges(): next_node = self.graph.graph.vs[out_edge.target] while True: if next_node['node_type'] == ExtendedOperator.QUANTIZE: if pre_quantize is None: pre_quantize = next_node['op'].outputs[0].quantization else: cur_quantize = next_node['op'].outputs[0].quantization if ( pre_quantize.scale != cur_quantize.scale or pre_quantize.zero_point != cur_quantize.zero_point or pre_quantize.dim != cur_quantize.dim ): skip = True break elif next_node['node_type'] == ExtendedOperator.DEQUANTIZE: break elif next_node['node_type'] in (ExtendedOperator.RESHAPE, ExtendedOperator.TRANSPOSE): if next_node.outdegree() > 1: skip = True break else: next_node = self.graph.graph.vs[next_node.out_edges()[0].target] else: skip = True break if skip: break if skip: continue # Find out the output of the first node in the sequence output_name = activ['op'].inputs[0].name output_idx = pre_activ['outputs'].index(output_name) new_output = pre_activ['outputs'][output_idx] assert new_output in self.graph.tensor_map # For each node that is next of the last node, we connect it with the first node # Also, the replace the tensors when needed self.graph.replace_next_tensors(activ, pre_activ, new_output) new_tensor = pre_activ['op'].outputs[0] old_tensor = activ['op'].outputs[0] new_tensor.quantization = old_tensor.quantization else: # Find out the output of the batch-norm nodes new_output = activ['outputs'][0] assert new_output in self.graph.tensor_map # For each node that is next of the activation node, we connect it with the previous node self.graph.connect_next_tensors(activ, pre_activ, new_output) # Update graph, prepare to drop the output tensor of the conv node and use the output tensor of the # batch-norm instead pre_activ['outputs'][0] = new_output pre_activ['op'].outputs[0] = self.graph.tensor_map[new_output] self.graph.tensor_node_map[new_output] = pre_activ['name'] tensor['name'] = activ['outputs'][0] tensor['label'] = activ['outputs'][0] remove_ids.append(activ.index) # Delete activation nodes self.graph.graph.delete_vertices(remove_ids) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_reciprocal_sqrt(self): # Find fusable ops edges = self.graph.graph.es.select(functools.partial(is_reciprocal_sqrt_edge, graph_converter=self.graph.graph)) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target], x) for x in edges) remove_ids = [] for sqrt, div, tensor in filtered_pairs: sqrt['node_type'] = ExtendedOperator.RSQRT sqrt['op'] = tfl.RsqrtOperator(sqrt['op'].inputs, sqrt['op'].outputs) div_op = div['op'] if ( div_op.inputs[0].buffer is not None and np.all(div_op.inputs[0].tensor == 1.0) and div['op'].fusedActivationFunction == ActivationFunctionType.NONE ): new_output = div['outputs'][0] assert new_output in self.graph.tensor_map # For each node that is next of the div node, we connect it with the previous node self.graph.connect_next_tensors(div, sqrt, new_output) # Update graph, prepare to drop the output tensor of the div node and use the output tensor of the # sqrt instead sqrt['outputs'][0] = new_output sqrt['op'].outputs[0] = self.graph.tensor_map[new_output] self.graph.tensor_node_map[new_output] = sqrt['name'] tensor['name'] = div['outputs'][0] tensor['label'] = div['outputs'][0] # remove div op remove_ids.append(div.index) else: div['node_type'] = ExtendedOperator.MUL div['op'] = tfl.MulOperator( div['op'].inputs, div['op'].outputs, fusedActivationFunction=div['op'].fusedActivationFunction ) # Delete div nodes self.graph.graph.delete_vertices(remove_ids) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def remove_tile_before_binary_elementwise_ops(self): # Find fusable ops edges = self.graph.graph.es.select(functools.partial(is_tile_binary_op_edge, graph_converter=self.graph.graph)) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target], x) for x in edges) remove_ids = [] actions = [] binary_op_ids = set() for tile, op_node, tensor in filtered_pairs: tile_op = tile['op'] binary_op = op_node['op'] input_idx = None for i in range(2): try: _ = tile['outputs'].index(binary_op.inputs[i].name) input_idx = i break except ValueError: pass if input_idx is None: continue alter_input_idx = 1 - input_idx try: out_shape = np.broadcast_shapes(binary_op.inputs[alter_input_idx].shape, tile_op.inputs[0].shape) if out_shape != binary_op.outputs[0].shape: continue except ValueError: continue if op_node.index not in binary_op_ids: binary_op_ids.add(op_node.index) else: continue new_tensor = tile_op.inputs[0] # Replace input tensors actions.append((self.graph.replace_operator_input, (op_node, input_idx, new_tensor))) # remove tile op remove_ids.append(tile.index) # Process actions for func, args in actions: func(*args) # Delete tile nodes self.graph.graph.delete_vertices(remove_ids) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_conv2d_gather(self): # Find fusable ops edges = self.graph.graph.es.select(functools.partial(is_conv2d_gather_edge, graph_converter=self.graph.graph)) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target], x) for x in edges) remove_ids = [] actions = [] for conv, gather, tensor in filtered_pairs: # Find out the output of the batch-norm nodes new_output = gather['outputs'][0] assert new_output in self.graph.tensor_map # For each node that is next of the activation node, we connect it with the previous node self.graph.connect_next_tensors(gather, conv, new_output) # Update graph, prepare to drop the output tensor of the gather node and use the output tensor of the # conv instead conv['outputs'][0] = new_output conv_out_quant_param = conv['op'].outputs[0].quantization conv['op'].outputs[0] = self.graph.tensor_map[new_output] conv['op'].outputs[0].quantization = conv_out_quant_param self.graph.tensor_node_map[new_output] = conv['name'] tensor['name'] = gather['outputs'][0] tensor['label'] = gather['outputs'][0] # permute weight of conv-op indx = gather['op'].inputs[1].tensor.copy() w = conv['op'].inputs[1].tensor.copy() w_quant_param = conv['op'].inputs[1].quantization new_w = np.take(w, indx, axis=0) # permute bias of conv-op b = conv['op'].inputs[2].tensor.copy() if len(conv['op'].inputs) > 2 else None b_quant_param = conv['op'].inputs[2].quantization new_b = np.take(b, indx, axis=0) if b is not None else None if w_quant_param is not None and isinstance(w_quant_param.scale, list) and w_quant_param.dim == 0: new_w_scale = np.take(w_quant_param.scale, indx, axis=0) new_w_zeros = np.take(w_quant_param.zero_point, indx, axis=0) w_quant_param.scale = new_w_scale w_quant_param.zero_point = new_w_zeros if new_b is not None: new_b_scale = np.take(b_quant_param.scale, indx, axis=0) new_b_zeros = np.take(b_quant_param.zero_point, indx, axis=0) b_quant_param.scale = new_b_scale b_quant_param.zero_point = new_b_zeros new_w = self.create_attr_tensor(new_w, quantization=w_quant_param) actions.append((self.graph.replace_operator_input, (conv, 1, new_w))) new_b = self.create_attr_tensor(new_b, quantization=b_quant_param) actions.append((self.graph.replace_operator_input, (conv, 2, new_b))) # remove gather op remove_ids.append(gather.index) # Process actions for func, args in actions: func(*args) # Delete activation nodes self.graph.graph.delete_vertices(remove_ids) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_gather_conv2d(self): # Find fusable ops edges = self.graph.graph.es.select(functools.partial(is_gather_conv2d_edge, graph_converter=self.graph.graph)) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]) for x in edges) def _remove_last_pred(seq): gather = seq[0] conv = seq[1] return False, (gather, conv) def _remove_last_action(first_node, last_node, custom_data): gather, conv = custom_data actions = [] indx = np.argsort(gather['op'].inputs[1].tensor) w = conv['op'].inputs[1].tensor.copy() w_quant_param = conv['op'].inputs[1].quantization new_w = np.take(w, indx, axis=3) if w_quant_param is not None and isinstance(w_quant_param.scale, list) and w_quant_param.dim == 3: new_w_scale = np.take(w_quant_param.scale, indx, axis=0) new_w_zeros = np.take(w_quant_param.zero_point, indx, axis=0) w_quant_param.scale = new_w_scale w_quant_param.zero_point = new_w_zeros new_w = self.create_attr_tensor(new_w, quantization=w_quant_param) actions.append((self.graph.replace_operator_input, (conv, 1, new_w))) return actions elinimate_sequences( self.graph, filtered_pairs, True, None, _remove_last_pred, _remove_last_action, False, force_forward_input=True, ) @class_conditional(lambda self: self.tflite_micro_rewrite) def split_requantize(self): vertices = self.graph.graph.vs.select(functools.partial(is_requantize_node, graph_converter=self.graph.graph)) remove_ids = [] ops = [] restore_mapping = [] for quantize in vertices: restore_nodes = [] # For each node that is next of a transformable node, # a. if it is an output node, remove it anyway since it will always be reconstructed # b. otherwise, record the info of the edge so that we may restore it after reconstruction for out_edge in quantize.out_edges(): next_node = self.graph.graph.vs[out_edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: remove_ids.append(next_node.index) del self.graph.tensor_map[next_node['outputs'][0]] del self.graph.tensor_node_map[next_node['outputs'][0]] else: restore_nodes.append((out_edge['name'], next_node['name'])) # Remove the mapping since they are going to be removed for output_name in quantize['outputs']: del self.graph.tensor_map[output_name] del self.graph.tensor_node_map[output_name] restore_mapping.append(restore_nodes) remove_ids.append(quantize.index) # Make sure the nodes are topologically sorted sorted_ops = [node['op'] for node in sorted(vertices, key=lambda x: int(re.search(r'\d+', x['name'])[0]))] # Delete nodes before transformation in the graph self.graph.graph.delete_vertices(remove_ids) for quantize, mapping in zip(sorted_ops, restore_mapping): input_tensor = quantize.inputs[0] output_tensor = quantize.outputs[0] intermediate = self.create_transform_tensor(input_tensor.tensor.astype('float32')) ops.append(tfl.DequantizeOperator([input_tensor], [intermediate])) ops.append(tfl.QuantizeOperator([intermediate], [output_tensor])) for op in ops: self.graph.add_operator(op, transform=True) self.graph.try_restore_edges(mapping) def transform_graph(self): # Find transformable ops filtered_nodes = self.graph.graph.vs.select( functools.partial(is_transformable_node, graph_converter=self.graph.graph) ) remove_ids = [] ops = [] restore_mapping = [] for node in filtered_nodes: restore_nodes = [] # For each node that is next of a transformable node, # a. if it is an output node, remove it anyway since it will always be reconstructed # b. otherwise, record the info of the edge so that we may restore it after reconstruction for out_edge in node.out_edges(): next_node = self.graph.graph.vs[out_edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: remove_ids.append(next_node.index) del self.graph.tensor_map[next_node['outputs'][0]] del self.graph.tensor_node_map[next_node['outputs'][0]] else: restore_nodes.append((out_edge['name'], next_node['name'])) # Remove the mapping since they are going to be removed for output_name in node['outputs']: del self.graph.tensor_map[output_name] del self.graph.tensor_node_map[output_name] restore_mapping.append(restore_nodes) ops.append(node) remove_ids.append(node.index) # Make sure the nodes are topologically sorted sorted_ops = [node['op'] for node in sorted(ops, key=lambda x: int(re.search(r'\d+', x['name'])[0]))] # Delete nodes before transformation in the graph self.graph.graph.delete_vertices(remove_ids) # Do transformation for op, mapping in zip(sorted_ops, restore_mapping): op.transform(self.graph, mapping) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_simple_transpose_pass(self): edges = self.graph.graph.es.select( functools.partial(is_transpose_fusable_edge, graph_converter=self.graph.graph) ) filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in edges] # Try to fuse the edges filtered_pairs = fuse_connected_edges(filtered_pairs) def _remove_first_pred(seq): new_perm = fuse_transpose_perms(seq) hints = set() for node in seq: if 'direction' in node['op'].extra_hints: hints.add(node['op'].extra_hints['direction']) if len(hints) == 1: hint = next(iter(hints)) else: hint = None remove_first = np.array_equal(new_perm, np.sort(new_perm)) return remove_first, (new_perm, hint) def _remove_first_action(first_node, last_node, custom_data): # Set fused perm to the first transpose node new_perm, hint = custom_data if hint is None: if 'direction' in first_node['op'].extra_hints: del first_node['op'].extra_hints['direction'] else: first_node['op'].extra_hints['direction'] = hint new_perm_tensor = self.create_attr_tensor(new_perm) action = (self.graph.replace_operator_input, (first_node, 1, new_perm_tensor)) return [action] elinimate_sequences(self.graph, filtered_pairs, _remove_first_pred, _remove_first_action) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_simple_gather_pass(self): edges = self.graph.graph.es.select(functools.partial(is_gather_fusable_edge, graph_converter=self.graph.graph)) filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in edges] # Try to fuse the edges filtered_pairs = fuse_connected_edges(filtered_pairs) def _remove_first_pred(seq): new_perm = fuse_transpose_perms(seq) hints = set() for node in seq: if 'direction' in node['op'].extra_hints: hints.add(node['op'].extra_hints['direction']) if len(hints) == 1: hint = next(iter(hints)) else: hint = None remove_first = np.array_equal(new_perm, np.sort(new_perm)) return remove_first, (new_perm, hint) def _remove_first_action(first_node, last_node, custom_data): # Set fused perm to the first transpose node new_perm, hint = custom_data if hint is None: if 'direction' in first_node['op'].extra_hints: del first_node['op'].extra_hints['direction'] else: first_node['op'].extra_hints['direction'] = hint new_perm_tensor = self.create_attr_tensor(new_perm) action = (self.graph.replace_operator_input, (first_node, 1, new_perm_tensor)) return [action] def _skip_pred(seq): for node in seq: op = node['op'] idx_tensor = op.inputs[1] if idx_tensor.buffer is None: return True if len(idx_tensor.shape) > 1: return True return False elinimate_sequences(self.graph, filtered_pairs, _remove_first_pred, _remove_first_action, skip_pred=_skip_pred) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_dequant_quant_pass(self, q_first): edges = self.graph.graph.es.select( functools.partial(is_dequant_quant_fusable_edge, graph_converter=self.graph.graph, q_first=q_first) ) filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in edges] r_edges = self.graph.graph.es.select( functools.partial(is_dequant_quant_fusable_edge, graph_converter=self.graph.graph, q_first=not q_first) ) r_filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in r_edges] filtered_pairs = fuse_connected_edges(filtered_pairs + r_filtered_pairs) new_pairs = [] for seq in filtered_pairs: start_idx = 0 end_idx = len(seq) if q_first: if seq[0]['node_type'] != ExtendedOperator.QUANTIZE: start_idx += 1 if seq[-1]['node_type'] != ExtendedOperator.DEQUANTIZE: end_idx -= 1 else: if seq[0]['node_type'] != ExtendedOperator.DEQUANTIZE: start_idx += 1 if seq[-1]['node_type'] != ExtendedOperator.QUANTIZE: end_idx -= 1 new_seq = seq[start_idx:end_idx] if len(new_seq) >= 2: new_pairs.append(new_seq) filtered_pairs = new_pairs def _remove_first_pred(seq): first_node, last_node = seq[0], seq[-1] new_qparams = last_node['op'].outputs[0].quantization orig_qparams = first_node['op'].inputs[0].quantization if ( first_node['node_type'] == ExtendedOperator.DEQUANTIZE and last_node['node_type'] == ExtendedOperator.QUANTIZE ): assert new_qparams is not None assert orig_qparams is not None remove_first = ( new_qparams.scale == orig_qparams.scale and new_qparams.zero_point == orig_qparams.zero_point and new_qparams.dim == orig_qparams.dim ) else: assert new_qparams is None assert orig_qparams is None remove_first = True return remove_first, None def _remove_first_action(first_node, last_node, custom_data): # Set new node type to first node first_node['node_type'] = ExtendedOperator.QUANTIZE old_op = first_node['op'] first_node['op'] = tfl.QuantizeOperator(old_op.inputs, old_op.outputs) return [] elinimate_sequences(self.graph, filtered_pairs, _remove_first_pred, _remove_first_action) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_simple_reshape_pass(self): edges = self.graph.graph.es.select(functools.partial(is_reshape_fusable_edge, graph_converter=self.graph.graph)) filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in edges] # Try to fuse the edge filtered_pairs = fuse_connected_edges(filtered_pairs) def _remove_first_pred(seq): first_node, last_node = seq[0], seq[-1] new_shape = last_node['op'].inputs[1].tensor orig_shape = np.array(first_node['op'].inputs[0].shape, dtype='int32') hints = set() for node in seq: if 'direction' in node['op'].extra_hints: hints.add(node['op'].extra_hints['direction']) if len(hints) == 1: hint = next(iter(hints)) else: hint = None remove_first = np.array_equal(new_shape, orig_shape) return remove_first, (new_shape, hint) def _remove_first_action(first_node, last_node, custom_data): # Set final shape to the first reshape node new_shape, hint = custom_data if hint is None: if 'direction' in first_node['op'].extra_hints: del first_node['op'].extra_hints['direction'] else: first_node['op'].extra_hints['direction'] = hint new_shape_tensor = self.create_attr_tensor(np.array(new_shape, dtype='int32')) first_node['op'].newShape = new_shape_tensor.tensor action = (self.graph.replace_operator_input, (first_node, 1, new_shape_tensor)) return [action] elinimate_sequences(self.graph, filtered_pairs, _remove_first_pred, _remove_first_action) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_simple_slice_pass(self): edges = self.graph.graph.es.select(functools.partial(is_slice_fusable_edge, graph_converter=self.graph.graph)) filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in edges] # Try to fuse the edge filtered_pairs = fuse_connected_edges(filtered_pairs) def _remove_first_pred(seq): fused_info = fuse_slices(seq) return False, fused_info def _remove_first_action(first_node, last_node, custom_data): # Set final shape to the first reshape node start, end, stride = custom_data if all((x == 1 for x in stride)): target_class = tfl.SliceOperator target_type = ExtendedOperator.SLICE else: target_class = tfl.StridedSliceOperator target_type = ExtendedOperator.STRIDED_SLICE if target_type == ExtendedOperator.SLICE: size = end - start start_tensor = self.create_attr_tensor(np.array(start, dtype='int32')) size_tensor = self.create_attr_tensor(np.array(size, dtype='int32')) actions = [ (self.graph.replace_operator_input, (first_node, 1, start_tensor)), (self.graph.replace_operator_input, (first_node, 2, size_tensor)), ] if first_node['node_type'] != ExtendedOperator.SLICE: old_slice_op = first_node['op'] first_node['node_type'] = ExtendedOperator.SLICE first_node['op'] = target_class(old_slice_op.inputs, old_slice_op.outputs) actions.append((self.graph.remove_operator_input, (first_node, 3))) else: size = end - start start_tensor = self.create_attr_tensor(np.array(start, dtype='int32')) end_tensor = self.create_attr_tensor(np.array(end, dtype='int32')) stride_tensor = self.create_attr_tensor(np.array(stride, dtype='int32')) if first_node['node_type'] == ExtendedOperator.STRIDED_SLICE: actions = [ (self.graph.replace_operator_input, (first_node, 1, start_tensor)), (self.graph.replace_operator_input, (first_node, 2, end_tensor)), (self.graph.replace_operator_input, (first_node, 3, stride_tensor)), ] else: old_slice_op = first_node['op'] first_node['node_type'] = ExtendedOperator.STRIDED_SLICE first_node['op'] = target_class(old_slice_op.inputs, old_slice_op.outputs) actions = [ (self.graph.replace_operator_input, (first_node, 1, start_tensor)), (self.graph.replace_operator_input, (first_node, 2, end_tensor)), (self.graph.append_operator_input, (first_node, stride_tensor)), ] return actions elinimate_sequences(self.graph, filtered_pairs, _remove_first_pred, _remove_first_action) @class_conditional(lambda self: self.group_tensors) def group_tensors_pass(self): tensor_map = {} actions = [] bytes_saved = 0 tensors_saved = 0 for v in self.graph.graph.vs: if v['node_type'] == ExtendedOperator.CONSTANT_NODE: tensor = self.graph.tensor_map[v['outputs'][0]] if tensor.quantization is None: t_idx = (tensor.buffer.data, tensor.dtype, tensor.shape) else: scale = tensor.quantization.scale zero_point = tensor.quantization.zero_point if isinstance(scale, list): scale = tuple(scale) if isinstance(zero_point, list): zero_point = tuple(zero_point) t_idx = ( tensor.buffer.data, tensor.dtype, tensor.shape, scale, zero_point, tensor.quantization.dim, ) if t_idx in tensor_map: new_tensor = tensor_map[t_idx] for e in v.out_edges(): target = e.target_vertex if target['op'] is not None: for i, inp in enumerate(target['op'].inputs): if inp.name == tensor.name: log.debug(f'{inp.name} used in {target["outputs"][0]}:{i} -> {new_tensor.name}') tensors_saved += 1 bytes_saved += len(inp.buffer.data) actions.append((self.graph.replace_operator_input, (target, i, new_tensor))) else: tensor_map[t_idx] = tensor # Process actions for func, args in actions: func(*args) log.info(f'{tensors_saved} duplicated tensors found, {bytes_saved / 1024 / 1024:.2f} MB saved') def cleanup_dead_nodes(self): cleanup_nodes = [] if not self.graph.graph.is_connected('weak'): while True: for vertex in self.graph.graph.vs: if ( vertex['node_type'] not in (ExtendedOperator.OUTPUT_NODE, ExtendedOperator.UNUSED_NODE) and vertex.outdegree() == 0 ): if vertex['node_type'] == ExtendedOperator.INPUT_NODE: continue if vertex['node_type'] != ExtendedOperator.CONSTANT_NODE: if vertex['op'] is None or vertex['op'].extra_hints.get('warn_on_unused', True): warnings.warn('Non constant node removed, something must be wrong there') log.warning('-' * 30) log.warning('Info of the deleted node:') log.warning(f'vertex: {vertex}') # edge = self.graph.graph.es.select(name=vertex['outputs'][0]) # assert edge is None, ( # f'The edge {vertex["outputs"][0]} exists but the connection to the vertex' # f' {vertex["name"]} is broken, probably there have some conflicts in the names' # ' of the nodes' # ) cleanup_nodes.append(vertex.index) if len(cleanup_nodes) == 0: break self.graph.graph.delete_vertices(cleanup_nodes) cleanup_nodes.clear() @class_conditional(lambda self: self.level >= GraphOptimizer.FOLD_BUFFER) def fold_transpose_buffer(self): edges = self.graph.graph.es.select( functools.partial(is_constant_transpose_fusable_edge, graph_converter=self.graph.graph) ) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target], x) for x in edges) remove_ids = [] for constant, transpose, tensor in filtered_pairs: # Calculate the output of the transposed constant nodes constant_tensor = transpose['op'].inputs[0].tensor perm_tensor = transpose['op'].inputs[1].tensor new_constant = np.transpose(constant_tensor, perm_tensor) new_tensor = self.create_attr_tensor(new_constant, quantization=transpose['op'].outputs[0].quantization) new_node = self.graph.add_nodes([new_tensor])[0] # For each node that is next of a constant transpose node, we connect it with the new constant node for out_edge in transpose.out_edges(): next_node = self.graph.graph.vs[out_edge.target] self.graph.graph.add_edge(new_node, next_node, name=new_tensor.name, label=new_tensor.name) log.debug( f'NEW EDGE: {new_node["label"]} -> {next_node["label"]} {self.graph.tensor_map[out_edge["name"]]}' ) op = next_node['op'] for idx in range(len(op.inputs)): if op.inputs[idx].name == transpose['op'].outputs[0].name: op.inputs[idx] = new_tensor remove_ids.append(transpose.index) # Delete constant transpose nodes self.graph.graph.delete_vertices(remove_ids) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def transpose_to_reshape_pass(self): filtered_nodes = self.graph.graph.vs.select( functools.partial(is_transformable_transpose_node, graph_converter=self.graph.graph) ) # Collect actions for the transformable transpose nodes actions = [] for node in filtered_nodes: original_op = node['op'] output_shape = np.array(original_op.outputs[0].shape, dtype='int32') shape_tensor = self.create_attr_tensor(output_shape) new_op = tfl.ReshapeOperator(original_op.inputs, original_op.outputs, output_shape) node['op'] = new_op node['node_type'] = ExtendedOperator.RESHAPE node['label'] = new_op.type_name() actions.append((self.graph.replace_operator_input, (node, 1, shape_tensor))) # Process actions for func, args in actions: node = args[0] func(*args) @class_conditional(lambda self: self.level >= GraphOptimizer.FOLD_BUFFER) def fold_reshape_buffer(self): edges = self.graph.graph.es.select( functools.partial(is_constant_reshape_fusable_edge, graph_converter=self.graph.graph) ) filtered_pairs = ((self.graph.graph.vs[x.source], self.graph.graph.vs[x.target], x) for x in edges) remove_ids = [] for constant, reshape, tensor in filtered_pairs: # Calculate the output of the transposed constant nodes constant_tensor = reshape['op'].inputs[0].tensor shape_tensor = reshape['op'].inputs[1].tensor new_constant = np.reshape(constant_tensor, shape_tensor) new_tensor = self.create_attr_tensor(new_constant, quantization=reshape['op'].inputs[0].quantization) new_node = self.graph.add_nodes([new_tensor])[0] # For each node that is next of a constant transpose node, we connect it with the new constant node for out_edge in reshape.out_edges(): next_node = self.graph.graph.vs[out_edge.target] self.graph.graph.add_edge(new_node, next_node, name=new_tensor.name, label=new_tensor.name) log.debug( f'NEW EDGE: {new_node["label"]} -> {next_node["label"]} {self.graph.tensor_map[out_edge["name"]]}' ) op = next_node['op'] for idx in range(len(op.inputs)): if op.inputs[idx].name == reshape['op'].outputs[0].name: op.inputs[idx] = new_tensor remove_ids.append(reshape.index) # Delete constant transpose nodes self.graph.graph.delete_vertices(remove_ids) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def remove_noop_pass(self, branch: bool = False): edges = self.graph.graph.es.select( functools.partial(is_ending_with_noop_edge, graph_converter=self.graph.graph, branch=branch) ) filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in edges] # Try to fuse the edges if not branch: filtered_pairs = fuse_connected_edges(filtered_pairs) elinimate_sequences(self.graph, filtered_pairs) @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_wrapped_reshape_within_transpose_pass(self): edges = self.graph.graph.es.select( functools.partial(is_wrapped_reshape_within_transpose_edge, graph_converter=self.graph.graph) ) filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in edges] # Try to fuse the edges fused_pairs = fuse_connected_edges(filtered_pairs) # Only TRANSPOSE->RESHAPE->TRANSPOSE is supported here filtered_pairs = [] for seq in fused_pairs: seq_len = len(seq) transpose_first = seq[0]['node_type'] == ExtendedOperator.TRANSPOSE if seq_len >= 3 and transpose_first: filtered_pairs.append(seq[:3]) elif seq_len >= 4: filtered_pairs.append(seq[1:4]) def _skip_pred(seq): mid_node = seq[1] orig_shape = mid_node['op'].inputs[0].shape new_shape = mid_node['op'].outputs[0].shape if not is_simple_reshape(orig_shape, new_shape): return True new_perm = fuse_transpose_perms_extended(seq) return (new_perm != np.sort(new_perm)).any() def _remove_last_pred(seq): orig_tensor = seq[0]['op'].inputs[0].tensor return False, (seq[2], orig_tensor) def _remove_last_action(first_node, last_node, custom_data): # Set final shape to the first reshape node last_trans, orig_tensor = custom_data actions = [] original_op = last_trans['op'] output_shape = np.array(original_op.outputs[0].shape, dtype='int32') shape_tensor = self.create_attr_tensor(output_shape) new_op = tfl.ReshapeOperator(original_op.inputs, original_op.outputs, output_shape) last_trans['op'] = new_op last_trans['node_type'] = ExtendedOperator.RESHAPE last_trans['label'] = new_op.type_name() new_op.inputs[0].tensor = orig_tensor new_op.inputs[0].shape = new_op.inputs[0].tensor.shape actions.append((self.graph.replace_operator_input, (last_trans, 1, shape_tensor))) return actions elinimate_sequences(self.graph, filtered_pairs, True, None, _remove_last_pred, _remove_last_action, _skip_pred) @class_conditional(lambda self: self.level >= GraphOptimizer.BRANCH_OPTIMIZE) def branch_reshape_expand_pass(self): edges = self.graph.graph.es.select(functools.partial(is_reshape_branch_edge, graph_converter=self.graph.graph)) branch_reshape_nodes = list(set(self.graph.graph.vs[edge.source] for edge in edges)) def _new_reshape(node: ig.Vertex, prev_node: ig.Vertex, next_node: ig.Vertex): actions = [] op = node['op'] op_out = op.outputs[0] op_shape = op.inputs[1] prev_idx = prev_node['outputs'].index(op.inputs[0].name) if prev_node['node_type'] == ExtendedOperator.INPUT_NODE: prev_out = self.graph.tensor_map[op.inputs[0].name] else: prev_op = prev_node['op'] prev_out = prev_op.outputs[prev_idx] new_tensor = self.create_transform_tensor(op_out.tensor.copy(), quantization=op_out.quantization) new_shape = self.create_attr_tensor(op_shape.tensor.copy()) new_op = tfl.ReshapeOperator([prev_out, new_shape], [new_tensor], new_shape.tensor) new_op.extra_hints.update(op.extra_hints) self.graph.add_operator(new_op) next_indices = [] for i, t in enumerate(next_node['op'].inputs): if t.name == op_out.name: actions.append((self.graph.replace_operator_input, (next_node, i, new_tensor))) next_indices.append(i) assert len(next_indices) > 0, f'{op_out.name} not in {[t.name for t in next_node["op"].inputs]}' return actions expand_op_outputs_in_branches(branch_reshape_nodes, _new_reshape, self.graph) @class_conditional(lambda self: self.level >= GraphOptimizer.BRANCH_OPTIMIZE) def branch_transpose_expand_pass(self): edges = self.graph.graph.es.select( functools.partial(is_transpose_branch_edge, graph_converter=self.graph.graph) ) branch_transpose_nodes = list(set(self.graph.graph.vs[edge.source] for edge in edges)) def _new_transpose(node: ig.Vertex, prev_node: ig.Vertex, next_node: ig.Vertex): actions = [] op = node['op'] op_out = op.outputs[0] op_perm = op.inputs[1] prev_idx = prev_node['outputs'].index(op.inputs[0].name) if prev_node['node_type'] in (ExtendedOperator.INPUT_NODE, ExtendedOperator.CONSTANT_NODE): prev_out = self.graph.tensor_map[op.inputs[0].name] else: prev_op = prev_node['op'] prev_out = prev_op.outputs[prev_idx] new_tensor = self.create_transform_tensor(op_out.tensor.copy(), quantization=op_out.quantization) new_perm = self.create_attr_tensor(op_perm.tensor.copy()) new_op = tfl.TransposeOperator([prev_out, new_perm], [new_tensor]) new_op.extra_hints.update(op.extra_hints) self.graph.add_operator(new_op) next_indices = [] for i, t in enumerate(next_node['op'].inputs): if t.name == op_out.name: actions.append((self.graph.replace_operator_input, (next_node, i, new_tensor))) next_indices.append(i) assert len(next_indices) > 0, f'{op_out.name} not in {[t.name for t in next_node["op"].inputs]}' return actions expand_op_outputs_in_branches(branch_transpose_nodes, _new_transpose, self.graph) @class_conditional(lambda self: self.level >= GraphOptimizer.BRANCH_OPTIMIZE, 0) def elementwise_reshape_transpose_passthrough_pass(self) -> int: edges = self.graph.graph.es.select( functools.partial(is_transpose_reshape_op_edge, graph_converter=self.graph.graph) ) pairs = ((self.graph.graph.vs[edge.source], self.graph.graph.vs[edge.target]) for edge in edges) filtered_nodes = (k[0] if k[0]['node_type'] != ExtendedOperator.TRANSPOSE else k[1] for k in pairs) unique_nodes = list(set(filtered_nodes)) actions = [] remove_edges = [] remove_vertices = [] processed_nodes = set() num_actions = 0 for node in unique_nodes: pending_processed_nodes = set() op = node['op'] input_indices = op_input_indices(op) l_shape = op.inputs[0].shape r_shape = op.outputs[0].shape if len(l_shape) == 0 or len(r_shape) == 0: continue l_map, r_map, _, _ = reshape_mapping(l_shape, r_shape) mode = None need_chain = False for l_val, r_val in zip(l_map, r_map): if len(l_val) > 1 and len(r_val) == 1: if mode in (None, 'up'): mode = 'up' else: mode = '?' break elif len(r_val) > 1 and len(l_val) == 1: if mode in (None, 'down'): mode = 'down' else: mode = '?' break elif len(r_val) > 1 and len(l_val) > 1: if len(r_val) != len(l_val) or r_val != l_val: # TODO: Support this case mode = '?' break else: need_chain = True if mode is None: mode = 'down' # TODO: Support multi-multi mappings if mode == '?': # reset hints if passthrough is not possible for i in input_indices: prev_node_name = op.inputs[i].name prev_node = self.graph.graph.vs.find(name=self.graph.tensor_node_map[prev_node_name]) if prev_node['node_type'] == ExtendedOperator.TRANSPOSE: if 'direction' in prev_node['op'].extra_hints: prev_node['op'].extra_hints.pop('direction') for edge in node.out_edges(): if edge.index in remove_edges: continue next_node = self.graph.graph.vs[edge.target] if 'direction' in next_node['op'].extra_hints: next_node['op'].extra_hints.pop('direction') continue check_consecutive_indices = [] if need_chain: new_l_map = [] new_r_map = [] for l_val, r_val in zip(l_map, r_map): if len(l_val) > 1 and len(r_val) > 1: if mode == 'down': check_consecutive_indices.append(l_val) else: check_consecutive_indices.append(r_val) for l_item in l_val: new_l_map.append([l_item]) for r_item in r_val: new_r_map.append([r_item]) else: new_l_map.append(l_val) new_r_map.append(r_val) l_map = new_l_map r_map = new_r_map prev_nodes = [] cand_perms = dict() cand_rev_perms = dict() prev_output_indices = [] num_constant_nodes = 0 prev_hints = set() skip = False for i in input_indices: prev_node_name = op.inputs[i].name prev_node = self.graph.graph.vs.find(name=self.graph.tensor_node_map[prev_node_name]) prev_nodes.append(prev_node) prev_output_indices.append(prev_node['outputs'].index(prev_node_name)) if prev_node['node_type'] == ExtendedOperator.TRANSPOSE: if prev_node['name'] in processed_nodes: skip = True break pending_processed_nodes.add(prev_node['name']) if mode == 'down': perm = tuple(prev_node['op'].inputs[1].tensor.tolist()) cand_perms.setdefault(perm, 0) cand_perms[perm] += 1 elif mode == 'up': perm = tuple(np.argsort(prev_node['op'].inputs[1].tensor).tolist()) cand_rev_perms.setdefault(perm, 0) cand_rev_perms[perm] += 1 if 'direction' in prev_node['op'].extra_hints: prev_hints.add(prev_node['op'].extra_hints['direction']) if prev_node['node_type'] == ExtendedOperator.CONSTANT_NODE: num_constant_nodes += 1 if skip or (self.level >= GraphOptimizer.BRANCH_OPTIMIZE_EXTENDED and 'up' in prev_hints): continue next_nodes = [] next_edges = [] out_nodes = [] next_hints = set() for edge in node.out_edges(): if edge.index in remove_edges: continue next_node = self.graph.graph.vs[edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: out_nodes.append(next_node) else: if next_node['name'] in processed_nodes: skip = True break pending_processed_nodes.add(next_node['name']) next_nodes.append(next_node) next_edges.append(edge) if next_node['node_type'] == ExtendedOperator.TRANSPOSE: if mode == 'down': perm = tuple(np.argsort(next_node['op'].inputs[1].tensor).tolist()) cand_rev_perms.setdefault(perm, 0) cand_rev_perms[perm] += 1 elif mode == 'up': perm = tuple(next_node['op'].inputs[1].tensor.tolist()) cand_perms.setdefault(perm, 0) cand_perms[perm] += 1 if 'direction' in next_node['op'].extra_hints: next_hints.add(next_node['op'].extra_hints['direction']) if skip or (self.level >= GraphOptimizer.BRANCH_OPTIMIZE_EXTENDED and 'down' in next_hints): continue cur_transpose_size = sum(cand_perms.values()) + sum(cand_rev_perms.values()) new_transpose_size = len(prev_nodes) + len(next_nodes) - sum(cand_perms.values()) - num_constant_nodes # Skip if the number of transpose nodes is not decreasing if len(cand_perms) == 0 or len(next_nodes) == 0 or new_transpose_size > cur_transpose_size: continue elif new_transpose_size == cur_transpose_size: skip = True if self.level >= GraphOptimizer.BRANCH_OPTIMIZE_EXTENDED: if 'down' in prev_hints or 'up' in next_hints: skip = False if skip: continue perm = max(cand_perms.items(), key=lambda x: x[1])[0] perm_arr = np.array(perm, dtype='int32') skip = False for check_idx in check_consecutive_indices: if mode == 'down': target_idx = perm_arr[check_idx] elif mode == 'up': perm_sorter = perm_arr.argsort() target_idx = perm_sorter[np.searchsorted(perm_arr, check_idx, sorter=perm_sorter)] normalized_src = [x - check_idx[0] for x in check_idx] normalized_tgt = [x - target_idx[0] for x in target_idx] if normalized_src != normalized_tgt: skip = True break if skip: continue num_actions += 1 remove_edges.extend([x.index for x in next_edges]) remove_vertices.extend([x.index for x in out_nodes]) for pending_processed_node in pending_processed_nodes: processed_nodes.add(pending_processed_node) for n in out_nodes: del self.graph.tensor_map[n['outputs'][0]] del self.graph.tensor_node_map[n['outputs'][0]] if mode == 'down': inv_perm_arr = np.argsort(perm_arr).astype('int32') l_dict = dict(zip([x[0] for x in l_map], r_map)) indices = map(lambda x: l_dict[x], inv_perm_arr.tolist()) inv_post_perm = list(itertools.chain.from_iterable(indices)) inv_post_perm_arr = np.array(inv_post_perm, dtype='int32') post_perm_arr = np.argsort(inv_post_perm_arr).astype('int32') elif mode == 'up': r_dict = dict(zip([x[0] for x in r_map], l_map)) indices = map(lambda x: r_dict[x], perm) inv_perm = list(itertools.chain.from_iterable(indices)) inv_perm_arr = np.array(inv_perm, dtype='int32') post_perm_arr = np.argsort(perm_arr).astype('int32') inv_post_perm_arr = np.argsort(post_perm_arr).astype('int32') for prev_node, prev_idx, next_idx in zip(prev_nodes, input_indices, prev_output_indices): if prev_node['op'] is None: prev_out = self.graph.tensor_map[prev_node['outputs'][0]] else: prev_out = prev_node['op'].outputs[next_idx] perm_tensor = self.create_attr_tensor(inv_perm_arr) prev_new_out = self.create_transform_tensor( np.transpose(prev_out.tensor, inv_perm_arr), quantization=prev_out.quantization ) transpose_op = tfl.TransposeOperator([prev_out, perm_tensor], [prev_new_out]) transpose_op.extra_hints['direction'] = 'up' self.graph.add_operator(transpose_op) actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True))) tensor_node_dict = {} for i, op_out in enumerate(op.outputs): perm_tensor = self.create_attr_tensor(post_perm_arr) new_out = self.create_transform_tensor( np.transpose(op_out.tensor, inv_post_perm_arr), quantization=op_out.quantization ) # Update relations if op_out.name in self.graph.tensor_node_map: del self.graph.tensor_node_map[op_out.name] self.graph.tensor_node_map[new_out.name] = node['name'] self.graph.tensor_map[new_out.name] = new_out node['outputs'][i] = new_out.name op.outputs[i] = new_out transpose_op = tfl.TransposeOperator([new_out, perm_tensor], [op_out]) transpose_op.extra_hints['direction'] = 'down' self.graph.add_operator(transpose_op) tensor_node_dict[op_out.name] = self.graph.graph.vs.find(name=self.graph.tensor_node_map[op_out.name]) # OP specific dim handling logic old_shape = op.inputs[1].tensor new_shape = self.create_attr_tensor(old_shape[inv_post_perm_arr]) actions.append((self.graph.replace_operator_input, (node, 1, new_shape, True))) op.newShape = new_shape.tensor for edge in next_edges: source = tensor_node_dict[edge['name']] self.graph.graph.add_edge(source, edge.target_vertex, name=edge['name'], label=edge['name']) # Process actions ids = [] for func, args in actions: node = args[0] res = func(*args) if res is not None: ids.extend(res) remove_edges = list(set(remove_edges + ids)) self.graph.graph.delete_edges(remove_edges) self.graph.graph.delete_vertices(remove_vertices) return num_actions @class_conditional(lambda self: self.rewrite_quantizable) def elementwise_op_quantize_passthrough_pass(self): edges = self.graph.graph.es.select( functools.partial( is_quantize_elementwise_op_edge, graph_converter=self.graph.graph, with_lstm=self.hybrid_int16_lstm ) ) pairs = ((self.graph.graph.vs[edge.source], self.graph.graph.vs[edge.target]) for edge in edges) filtered_nodes = (k[0] if k[0]['node_type'] != ExtendedOperator.DEQUANTIZE else k[1] for k in pairs) unique_nodes = list(set(filtered_nodes)) actions = [] remove_edges = [] remove_vertices = [] for node in unique_nodes: op = node['op'] input_indices = op_input_indices(op) prev_nodes = [] q_tensors = dict() prev_output_indices = [] skip_names = [] for i in input_indices: prev_node_name = op.inputs[i].name prev_node = self.graph.graph.vs.find(name=self.graph.tensor_node_map[prev_node_name]) prev_nodes.append(prev_node) prev_output_indices.append(prev_node['outputs'].index(prev_node_name)) if prev_node['node_type'] == ExtendedOperator.DEQUANTIZE: q_tensors[prev_node_name] = prev_node['op'].inputs[0] if prev_node['node_type'] == ExtendedOperator.CONSTANT_NODE: if ( node['node_type'] in (ExtendedOperator.MINIMUM, ExtendedOperator.MAXIMUM) and i != 0 and prev_node_name not in self.graph.q_mapping ): f_tensor = self.graph.tensor_map[prev_node_name] r_tensor = q_tensors[op.inputs[0].name] q_arr = np.rint( f_tensor.tensor / r_tensor.quantization.scale + r_tensor.quantization.zero_point ) i_type = np.iinfo(r_tensor.tensor.dtype) if np.any(q_arr > i_type.max): warnings.warn('Overflow while quantizing the tensor') q_arr = np.minimum(q_arr, i_type.max) if np.any(q_arr < i_type.min): warnings.warn('Underflow while quantizing the tensor') q_arr = np.maximum(q_arr, i_type.min) q_arr = q_arr.astype(r_tensor.dtype) q_tensor = self.create_attr_tensor(q_arr, quantization=r_tensor.quantization) self.graph.q_mapping[prev_node_name] = q_tensor if prev_node_name in self.graph.q_mapping: skip_names.append(prev_node_name) next_nodes = [] next_edges = [] out_nodes = [] for edge in node.out_edges(): if edge.index in remove_edges: continue next_node = self.graph.graph.vs[edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: out_nodes.append(next_node) else: next_nodes.append(next_node) next_edges.append(edge) if next_node['node_type'] == ExtendedOperator.QUANTIZE: skip = False name = next_node['op'].inputs[0].name q_tensor = next_node['op'].outputs[0] assert q_tensor.quantization is not None if node['node_type'] in ( ExtendedOperator.BATCH_MATMUL, ExtendedOperator.ABS, ExtendedOperator.RSQRT, ): if q_tensor.dtype not in (np.dtype('int8'), np.dtype('int16')): skip = True elif node['node_type'] == ExtendedOperator.DIV: if q_tensor.dtype != np.dtype('uint8'): skip = True elif node['node_type'] == ExtendedOperator.SOFTMAX: if q_tensor.dtype == np.dtype('int8'): if ( abs(q_tensor.quantization.scale - 1.0 / 256) > 0.001 * 1.0 / 256 or q_tensor.quantization.zero_point != -128 ): skip = True elif q_tensor.dtype == np.dtype('int16'): if ( abs(q_tensor.quantization.scale - 1.0 / 32768) > 0.001 * 1.0 / 32768 or q_tensor.quantization.zero_point != 0 ): skip = True elif q_tensor.dtype == np.dtype('uint8'): if ( abs(q_tensor.quantization.scale - 1.0 / 256) > 0.001 * 1.0 / 256 or q_tensor.quantization.zero_point != 0 ): log.warning( 'On some chips, only softmax with scale=1.0/256 and zero_point=0 is supported' ) else: skip = True elif node['node_type'] == ExtendedOperator.LOG_SOFTMAX: if q_tensor.dtype == np.dtype('int8'): if q_tensor.quantization.scale != 16.0 / 256 or q_tensor.quantization.zero_point != 127: skip = True elif q_tensor.dtype == np.dtype('uint8'): if q_tensor.quantization.scale != 16.0 / 256 or q_tensor.quantization.zero_point != 255: skip = True else: skip = True if not skip: q_tensors[name] = q_tensor cur_transpose_size = len(q_tensors) new_transpose_size = len(prev_nodes) + len(next_nodes) - len(skip_names) # Skip if the number of [de]quantize nodes is not decreasing if len(next_nodes) == 0 or new_transpose_size > cur_transpose_size: continue remove_edges.extend([x.index for x in next_edges]) remove_vertices.extend([x.index for x in out_nodes]) for n in out_nodes: del self.graph.tensor_map[n['outputs'][0]] del self.graph.tensor_node_map[n['outputs'][0]] tensor_node_dict = {} for prev_node, prev_idx, next_idx in zip(prev_nodes, input_indices, prev_output_indices): if prev_node['op'] is None: prev_out = self.graph.tensor_map[prev_node['outputs'][0]] else: prev_out = prev_node['op'].outputs[next_idx] if prev_out.name in tensor_node_dict: prev_new_out, skip = tensor_node_dict[prev_out.name] actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True, skip))) skip += 1 tensor_node_dict[prev_out.name] = (prev_new_out, skip) else: if prev_out.name in skip_names: prev_new_out = self.graph.q_mapping[prev_out.name] self.graph.add_nodes([prev_new_out]) tensor_node_dict[prev_out.name] = (prev_new_out, 1) actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True))) else: prev_new_out = self.create_transform_tensor( q_tensors[prev_out.name].tensor, quantization=q_tensors[prev_out.name].quantization ) tensor_node_dict[prev_out.name] = (prev_new_out, 1) self.graph.add_operator(tfl.QuantizeOperator([prev_out], [prev_new_out])) actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True))) tensor_node_dict = {} for i, op_out in enumerate(op.outputs): new_out = self.create_transform_tensor( q_tensors[op_out.name].tensor, quantization=q_tensors[op_out.name].quantization ) # Update relations if op_out.name in self.graph.tensor_node_map: del self.graph.tensor_node_map[op_out.name] self.graph.tensor_node_map[new_out.name] = node['name'] self.graph.tensor_map[new_out.name] = new_out node['outputs'][i] = new_out.name op.outputs[i] = new_out self.graph.add_operator(tfl.DequantizeOperator([new_out], [op_out])) tensor_node_dict[op_out.name] = self.graph.graph.vs.find(name=self.graph.tensor_node_map[op_out.name]) for edge in next_edges: source = tensor_node_dict[edge['name']] self.graph.graph.add_edge(source, edge.target_vertex, name=edge['name'], label=edge['name']) # Process actions ids = [] for func, args in actions: node = args[0] res = func(*args) if res is not None: ids.extend(res) remove_edges = list(set(remove_edges + ids)) self.graph.graph.delete_edges(remove_edges) self.graph.graph.delete_vertices(remove_vertices) @class_conditional(lambda self: self.level >= GraphOptimizer.BRANCH_OPTIMIZE, 0) def elementwise_op_transpose_passthrough_pass(self, quantizable_ops_only: bool = False) -> int: edges = self.graph.graph.es.select( functools.partial( is_transpose_elementwise_op_edge, graph_converter=self.graph.graph, quantizable_ops_only=quantizable_ops_only, ) ) pairs = ((self.graph.graph.vs[edge.source], self.graph.graph.vs[edge.target]) for edge in edges) if quantizable_ops_only: all_edges = self.graph.graph.es.select( functools.partial( is_transpose_elementwise_op_edge, graph_converter=self.graph.graph, quantizable_ops_only=False, ) ) all_pairs = ((self.graph.graph.vs[edge.source], self.graph.graph.vs[edge.target]) for edge in all_edges) forward_d = dict(all_pairs) backward_d = {v: k for k, v in forward_d.items()} filtered_nodes = [] for s, e in pairs: if s['node_type'] == ExtendedOperator.TRANSPOSE: pn = backward_d.get(s, None) if pn is not None: filtered_nodes.append(pn) else: log.warning('Cannot passthrough transpose upward around requantizable ops') else: pn = forward_d.get(e, None) if pn is not None: filtered_nodes.append(pn) else: log.warning('Cannot passthrough transpose downward around requantizable ops') else: filtered_nodes = (k[0] if k[0]['node_type'] != ExtendedOperator.TRANSPOSE else k[1] for k in pairs) unique_nodes = list(set(filtered_nodes)) actions = [] remove_edges = [] remove_vertices = [] num_actions = 0 for node in unique_nodes: op = node['op'] input_indices = op_input_indices(op) prev_nodes = [] cand_perms = dict() prev_output_indices = [] num_constant_nodes = 0 num_reshape_transpose = 0 prev_hints = set() for i in input_indices: prev_node_name = op.inputs[i].name prev_node = self.graph.graph.vs.find(name=self.graph.tensor_node_map[prev_node_name]) prev_nodes.append(prev_node) prev_output_indices.append(prev_node['outputs'].index(prev_node_name)) if prev_node['node_type'] == ExtendedOperator.TRANSPOSE: perm = tuple(prev_node['op'].inputs[1].tensor.tolist()) if node['node_type'] == ExtendedOperator.PACK: perm = [i if i < op.axis else i + 1 for i in perm] perm.insert(op.axis, op.axis) perm = tuple(perm) cand_perms.setdefault(perm, 0) cand_perms[perm] += 1 if 'direction' in prev_node['op'].extra_hints: prev_hints.add(prev_node['op'].extra_hints['direction']) if prev_node['node_type'] == ExtendedOperator.CONSTANT_NODE: num_constant_nodes += 1 if prev_node['node_type'] == ExtendedOperator.RESHAPE: prev_prev_node_name = self.graph.tensor_node_map[prev_node['op'].inputs[0].name] prev_prev_node = self.graph.graph.vs.find(name=prev_prev_node_name) if prev_prev_node['node_type'] == ExtendedOperator.TRANSPOSE: num_reshape_transpose += 1 if 'direction' in prev_prev_node['op'].extra_hints: prev_hints.add(prev_prev_node['op'].extra_hints['direction']) if self.level >= GraphOptimizer.BRANCH_OPTIMIZE_EXTENDED and 'up' in prev_hints: continue next_nodes = [] next_edges = [] out_nodes = [] skip_names = [] next_hints = set() for edge in node.out_edges(): if edge.index in remove_edges: continue next_node = self.graph.graph.vs[edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: out_nodes.append(next_node) elif next_node['node_type'] == ExtendedOperator.UNUSED_NODE: skip_names.append(edge['label']) else: next_nodes.append(next_node) next_edges.append(edge) if next_node['node_type'] == ExtendedOperator.TRANSPOSE: perm = tuple(np.argsort(next_node['op'].inputs[1].tensor).tolist()) if node['node_type'] == ExtendedOperator.UNPACK: perm = [i if i < op.axis else i + 1 for i in perm] perm.insert(op.axis, op.axis) perm = tuple(perm) cand_perms.setdefault(perm, 0) cand_perms[perm] += 1 if 'direction' in next_node['op'].extra_hints: next_hints.add(next_node['op'].extra_hints['direction']) if next_node['node_type'] == ExtendedOperator.RESHAPE: o_nodes = [e.target_vertex for e in next_node.out_edges()] if len(o_nodes) == 1 and o_nodes[0]['node_type'] == ExtendedOperator.TRANSPOSE: num_reshape_transpose += 1 if 'direction' in o_nodes[0]['op'].extra_hints: next_hints.add(o_nodes[0]['op'].extra_hints['direction']) if self.level >= GraphOptimizer.BRANCH_OPTIMIZE_EXTENDED and 'down' in next_hints: continue cur_transpose_size = sum(cand_perms.values()) + num_reshape_transpose new_transpose_size = ( len(prev_nodes) + len(next_nodes) - num_constant_nodes - cur_transpose_size + num_reshape_transpose ) # Skip if the following conditions are met # a. the number of transpose nodes is not decreasing (skip if `bypass_elementwise_passthrough_constraint`) # b. no hint can be found (skip if optimize level is below BRANCH_OPTIMIZE_EXTENDED) is_increasing = new_transpose_size > cur_transpose_size is_not_decreasing = new_transpose_size >= cur_transpose_size is_same = new_transpose_size == cur_transpose_size if len(next_nodes) == 0: continue else: if self.bypass_elementwise_passthrough_constraint: condition = is_not_decreasing else: if is_increasing: continue condition = is_same if condition: skip = True if self.level >= GraphOptimizer.BRANCH_OPTIMIZE_EXTENDED: if 'down' in prev_hints or 'up' in next_hints: skip = False if skip: continue num_actions += 1 remove_edges.extend([x.index for x in next_edges]) remove_vertices.extend([x.index for x in out_nodes]) for n in out_nodes: del self.graph.tensor_map[n['outputs'][0]] del self.graph.tensor_node_map[n['outputs'][0]] perm = max(cand_perms.items(), key=lambda x: x[1])[0] perm_arr = np.array(perm, dtype='int32') inv_perm_arr = np.argsort(perm_arr).astype('int32') if node['node_type'] == ExtendedOperator.UNPACK: inv_perm_arr_post = inv_perm_arr[inv_perm_arr != op.axis] inv_perm_arr_post[inv_perm_arr_post > op.axis] -= 1 perm_arr_post = np.argsort(inv_perm_arr_post).astype('int32') elif node['node_type'] == ExtendedOperator.PACK: perm_arr_post = perm_arr inv_perm_arr_post = inv_perm_arr perm_arr = perm_arr_post[perm_arr_post != op.axis] perm_arr[perm_arr > op.axis] -= 1 inv_perm_arr = np.argsort(perm_arr).astype('int32') else: perm_arr_post = perm_arr inv_perm_arr_post = inv_perm_arr tensor_node_dict = {} for prev_node, prev_idx, next_idx in zip(prev_nodes, input_indices, prev_output_indices): if prev_node['op'] is None: prev_out = self.graph.tensor_map[prev_node['outputs'][0]] else: prev_out = prev_node['op'].outputs[next_idx] if prev_out.name in tensor_node_dict: prev_new_out, skip = tensor_node_dict[prev_out.name] actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True, skip))) skip += 1 tensor_node_dict[prev_out.name] = (prev_new_out, skip) else: perm_tensor = self.create_attr_tensor(inv_perm_arr) if len(prev_out.shape) != perm_tensor.tensor.size: new_shape = [1] * (perm_tensor.tensor.size - len(prev_out.shape)) + list(prev_out.shape) prev_out_reshaped = self.create_transform_tensor( np.reshape(prev_out.tensor, new_shape), quantization=prev_out.quantization ) new_shape_tensor = self.create_attr_tensor(np.array(new_shape, dtype='int32')) self.graph.add_operator( tfl.ReshapeOperator([prev_out, new_shape_tensor], [prev_out_reshaped], new_shape) ) prev_out = prev_out_reshaped prev_new_out = self.create_transform_tensor( np.transpose(prev_out.tensor, inv_perm_arr), quantization=prev_out.quantization ) tensor_node_dict[prev_out.name] = (prev_new_out, 1) transpose_op = tfl.TransposeOperator([prev_out, perm_tensor], [prev_new_out]) transpose_op.extra_hints['direction'] = 'up' self.graph.add_operator(transpose_op) actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True))) tensor_node_dict = {} for i, op_out in enumerate(op.outputs): # For unused tensors, we perform inplace shape updates if op_out.name in skip_names: orig_shape = np.array(op_out.shape, dtype='int32') new_shape = orig_shape[inv_perm_arr] op_out.shape = tuple(new_shape.tolist()) continue perm_tensor = self.create_attr_tensor(perm_arr_post) new_out = self.create_transform_tensor( np.transpose(op_out.tensor, inv_perm_arr_post), quantization=op_out.quantization ) # Update relations if op_out.name in self.graph.tensor_node_map: del self.graph.tensor_node_map[op_out.name] self.graph.tensor_node_map[new_out.name] = node['name'] self.graph.tensor_map[new_out.name] = new_out node['outputs'][i] = new_out.name op.outputs[i] = new_out transpose_op = tfl.TransposeOperator([new_out, perm_tensor], [op_out]) transpose_op.extra_hints['direction'] = 'down' self.graph.add_operator(transpose_op) tensor_node_dict[op_out.name] = self.graph.graph.vs.find(name=self.graph.tensor_node_map[op_out.name]) # OP specific dim handling logic if node['node_type'] in (ExtendedOperator.CONCATENATION, ExtendedOperator.GATHER, ExtendedOperator.UNPACK): old_axis = op.axis new_axis = np.where(inv_perm_arr == old_axis)[0][0] op.axis = new_axis elif node['node_type'] == ExtendedOperator.PACK: old_axis = op.axis new_axis = np.where(inv_perm_arr_post == old_axis)[0][0] op.axis = new_axis elif node['node_type'] == ExtendedOperator.SPLIT_V: old_dim = op.inputs[2].tensor new_dim = np.where(inv_perm_arr == old_dim)[0][0] new_dim_tensor = self.create_attr_tensor(np.array(new_dim, dtype='int32')) actions.append((self.graph.replace_operator_input, (node, 2, new_dim_tensor, True))) elif node['node_type'] == ExtendedOperator.SPLIT: old_dim = op.inputs[0].tensor new_dim = np.where(inv_perm_arr == old_dim)[0][0] new_dim_tensor = self.create_attr_tensor(np.array(new_dim, dtype='int32')) actions.append((self.graph.replace_operator_input, (node, 0, new_dim_tensor, True))) elif node['node_type'] in ( ExtendedOperator.PAD, ExtendedOperator.PADV2, ExtendedOperator.MIRROR_PAD, ExtendedOperator.TILE, ): old_pad = op.inputs[1].tensor new_pad = self.create_attr_tensor(old_pad[inv_perm_arr]) actions.append((self.graph.replace_operator_input, (node, 1, new_pad, True))) elif node['node_type'] == ExtendedOperator.PRELU: old_weight = op.inputs[1].tensor if old_weight.ndim != 1: assert old_weight.ndim + 1 == len(inv_perm_arr) new_perm = np.argsort(np.argsort(inv_perm_arr[1:])) new_perm_t = self.create_attr_tensor(np.array(new_perm, dtype='int32')) new_weight = self.create_transform_tensor(np.transpose(old_weight, new_perm)) self.graph.add_operator(tfl.TransposeOperator([op.inputs[1], new_perm_t], [new_weight])) actions.append((self.graph.replace_operator_input, (node, 1, new_weight, True))) elif node['node_type'] in (ExtendedOperator.SLICE, ExtendedOperator.STRIDED_SLICE): for i, t in enumerate(op.inputs[1:]): if t.buffer is None: new_perm_t = self.create_attr_tensor(np.array(inv_perm_arr, dtype='int32')) new_t = self.create_transform_tensor(t.tensor[inv_perm_arr]) self.graph.add_operator(tfl.TransposeOperator([t, new_perm_t], [new_t])) else: new_t = self.create_attr_tensor(t.tensor[inv_perm_arr]) actions.append((self.graph.replace_operator_input, (node, i + 1, new_t, True))) elif node['node_type'] in ( ExtendedOperator.SUM, ExtendedOperator.ARG_MIN, ExtendedOperator.ARG_MAX, ExtendedOperator.REDUCE_MIN, ExtendedOperator.REDUCE_MAX, ExtendedOperator.REDUCE_PROD, ExtendedOperator.MEAN, ): old_axis = op.inputs[1].tensor.tolist() new_axis = [] for t in old_axis: new_t = np.where(inv_perm_arr == t)[0][0] new_axis.append(new_t) axis_arr = np.array(new_axis, dtype='int32') axis_tensor = self.create_attr_tensor(axis_arr) actions.append((self.graph.replace_operator_input, (node, 1, axis_tensor, True))) for edge in next_edges: source = tensor_node_dict[edge['name']] self.graph.graph.add_edge(source, edge.target_vertex, name=edge['name'], label=edge['name']) # Process actions ids = [] for func, args in actions: node = args[0] res = func(*args) if res is not None: ids.extend(res) remove_edges = list(set(remove_edges + ids)) self.graph.graph.delete_edges(remove_edges) self.graph.graph.delete_vertices(remove_vertices) return num_actions @class_conditional(lambda self: self.level >= GraphOptimizer.BRANCH_OPTIMIZE, 0) def elementwise_op_reshape_passthrough_pass(self) -> int: edges = self.graph.graph.es.select( functools.partial(is_reshape_elementwise_op_edge, graph_converter=self.graph.graph) ) pairs = ((self.graph.graph.vs[edge.source], self.graph.graph.vs[edge.target]) for edge in edges) filtered_nodes = (k[0] if k[0]['node_type'] != ExtendedOperator.RESHAPE else k[1] for k in pairs) unique_nodes = list(set(filtered_nodes)) actions = [] remove_edges = [] remove_vertices = [] num_actions = 0 for node in unique_nodes: op = node['op'] dim_indice = op_input_dims(op) input_indices = op_input_indices(op) prev_nodes = [] cand_shapes = dict() cand_next_shapes = dict() prev_output_indices = [] num_constant_nodes = 0 prev_hints = set() for i in input_indices: prev_node_name = op.inputs[i].name prev_node = self.graph.graph.vs.find(name=self.graph.tensor_node_map[prev_node_name]) prev_nodes.append(prev_node) prev_output_indices.append(prev_node['outputs'].index(prev_node_name)) if prev_node['node_type'] == ExtendedOperator.CONSTANT_NODE: num_constant_nodes += 1 if prev_node['node_type'] == ExtendedOperator.RESHAPE: mapping = dict() if not is_simple_reshape( prev_node['op'].inputs[0].shape, prev_node['op'].outputs[0].shape, mapping ): continue new_dim = None if dim_indice is not None: rev_mapping = {v: k for k, v in mapping.items()} if node['node_type'] == ExtendedOperator.PACK: if dim_indice in rev_mapping: tmp_new_dim = rev_mapping[dim_indice] else: if dim_indice - 1 in rev_mapping: tmp_new_dim = rev_mapping[dim_indice - 1] + 1 elif dim_indice + 1 in rev_mapping: tmp_new_dim = rev_mapping[dim_indice + 1] - 1 else: # TODO: Figure out the rev index tmp_new_dim = -1 tmp_dim_indice = dim_indice new_dim = -1 dim_indice = -1 else: if dim_indice not in rev_mapping: continue new_dim = rev_mapping[dim_indice] shape = tuple(prev_node['op'].inputs[0].shape) shape = tuple(x if i != new_dim else -1 for i, x in enumerate(shape)) if node['node_type'] == ExtendedOperator.PACK and tmp_new_dim >= 0: shape = list(shape) shape.insert(tmp_new_dim, -1) shape = tuple(shape) cand_shapes.setdefault(shape, 0) cand_shapes[shape] += 1 next_shape = tuple(prev_node['op'].outputs[0].shape) next_shape = tuple(x if i != dim_indice else -1 for i, x in enumerate(next_shape)) if node['node_type'] == ExtendedOperator.PACK: next_shape = list(next_shape) next_shape.insert(tmp_dim_indice, -1) next_shape = tuple(next_shape) cand_next_shapes.setdefault(next_shape, 0) cand_next_shapes[next_shape] += 1 if node['node_type'] == ExtendedOperator.PACK: dim_indice = tmp_dim_indice if 'direction' in prev_node['op'].extra_hints: prev_hints.add(prev_node['op'].extra_hints['direction']) if self.level >= GraphOptimizer.BRANCH_OPTIMIZE_EXTENDED and 'up' in prev_hints: continue next_nodes = [] next_edges = [] out_nodes = [] skip_names = [] next_hints = set() for edge in node.out_edges(): if edge.index in remove_edges: continue next_node = self.graph.graph.vs[edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: out_nodes.append(next_node) elif next_node['node_type'] == ExtendedOperator.UNUSED_NODE: skip_names.append(edge['label']) else: next_nodes.append(next_node) next_edges.append(edge) if next_node['node_type'] == ExtendedOperator.RESHAPE: mapping = dict() if not is_simple_reshape( next_node['op'].inputs[0].shape, next_node['op'].outputs[0].shape, mapping ): continue new_dim = None if dim_indice is not None: if node['node_type'] == ExtendedOperator.UNPACK: if dim_indice in mapping: tmp_new_dim = mapping[dim_indice] else: if dim_indice - 1 in mapping: tmp_new_dim = mapping[dim_indice - 1] + 1 elif dim_indice + 1 in mapping: tmp_new_dim = mapping[dim_indice + 1] - 1 else: # TODO: Figure out the rev index tmp_new_dim = -1 tmp_dim_indice = dim_indice new_dim = -1 dim_indice = -1 else: if dim_indice not in mapping: continue new_dim = mapping[dim_indice] shape = tuple(next_node['op'].outputs[0].shape) shape = tuple(x if i != new_dim else -1 for i, x in enumerate(shape)) if node['node_type'] == ExtendedOperator.UNPACK and tmp_new_dim >= 0: shape = list(shape) shape.insert(tmp_new_dim, -1) shape = tuple(shape) cand_shapes.setdefault(shape, 0) cand_shapes[shape] += 1 next_shape = tuple(next_node['op'].inputs[0].shape) next_shape = tuple(x if i != dim_indice else -1 for i, x in enumerate(next_shape)) if node['node_type'] == ExtendedOperator.UNPACK: next_shape = list(next_shape) next_shape.insert(tmp_dim_indice, -1) next_shape = tuple(next_shape) cand_next_shapes.setdefault(next_shape, 0) cand_next_shapes[next_shape] += 1 if node['node_type'] == ExtendedOperator.UNPACK: dim_indice = tmp_dim_indice if 'direction' in next_node['op'].extra_hints: next_hints.add(next_node['op'].extra_hints['direction']) if len(cand_shapes) == 0: continue if self.level >= GraphOptimizer.BRANCH_OPTIMIZE_EXTENDED and 'down' in next_hints: continue cur_reshape_size = max(cand_shapes.values()) cur_next_reshape_size = max(cand_next_shapes.values()) full_size = len(prev_nodes) + len(next_nodes) if cur_reshape_size != cur_next_reshape_size: continue new_reshape_size = full_size - cur_reshape_size - num_constant_nodes # Skip if not wrapped by reshapes if ( len(next_nodes) == 0 or new_reshape_size > cur_reshape_size ): # cur_reshape_size < full_size or cur_next_reshape_size < full_size: continue elif new_reshape_size == cur_reshape_size: skip = True if self.level >= GraphOptimizer.BRANCH_OPTIMIZE_EXTENDED: if 'down' in prev_hints or 'up' in next_hints: skip = False if skip: continue num_actions += 1 remove_edges.extend([x.index for x in next_edges]) remove_vertices.extend([x.index for x in out_nodes]) for n in out_nodes: del self.graph.tensor_map[n['outputs'][0]] del self.graph.tensor_node_map[n['outputs'][0]] prev_shape = max(cand_shapes.items(), key=lambda x: x[1])[0] next_shape = max(cand_next_shapes.items(), key=lambda x: x[1])[0] tensor_node_dict = {} for prev_node, prev_idx, next_idx in zip(prev_nodes, input_indices, prev_output_indices): if prev_node['op'] is None: prev_out = self.graph.tensor_map[prev_node['outputs'][0]] else: prev_out = prev_node['op'].outputs[next_idx] if prev_out.name in tensor_node_dict: prev_new_out, skip = tensor_node_dict[prev_out.name] actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True, skip))) skip += 1 tensor_node_dict[prev_out.name] = (prev_new_out, skip) else: if node['node_type'] == ExtendedOperator.PACK: tmp_prev_shape = prev_shape prev_shape = [i for i in prev_shape if i != -1] prev_shape_aligned = prev_shape if np.prod(prev_out.shape) != np.prod(prev_shape): new_prev_shape = prev_out.shape if len(prev_out.shape) < len(next_shape): new_prev_shape = [1] * (len(next_shape) - len(prev_out.shape)) + list(prev_out.shape) mapping = {} is_simple_reshape(prev_shape, next_shape, mapping) prev_shape_aligned = np.ones(len(prev_shape), dtype='int32') for pi, ni in mapping.items(): prev_shape_aligned[pi] = new_prev_shape[ni] prev_new_out = self.create_transform_tensor( np.reshape(prev_out.tensor, prev_shape_aligned), quantization=prev_out.quantization ) tensor_node_dict[prev_out.name] = (prev_new_out, 1) shape_tensor = self.create_attr_tensor(np.array(prev_new_out.shape, dtype='int32')) reshape_op = tfl.ReshapeOperator( [prev_out, shape_tensor], [prev_new_out], newShape=shape_tensor.tensor ) reshape_op.extra_hints['direction'] = 'up' self.graph.add_operator(reshape_op) actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True))) if node['node_type'] == ExtendedOperator.PACK: prev_shape = tmp_prev_shape tensor_node_dict = {} for i, op_out in enumerate(op.outputs): if node['node_type'] == ExtendedOperator.UNPACK: tmp_prev_shape = prev_shape prev_shape = [i for i in prev_shape if i != -1] # For unused tensors, we perform inplace shape updates if op_out.name in skip_names: new_shape = np.reshape(op_out.tensor, prev_shape).shape op_out.shape = tuple(new_shape) if node['node_type'] == ExtendedOperator.UNPACK: prev_shape = tmp_prev_shape continue new_out = self.create_transform_tensor( np.reshape(op_out.tensor, prev_shape), quantization=op_out.quantization ) shape_tensor = self.create_attr_tensor(np.array(op_out.shape, dtype='int32')) # Update relations if op_out.name in self.graph.tensor_node_map: del self.graph.tensor_node_map[op_out.name] self.graph.tensor_node_map[new_out.name] = node['name'] self.graph.tensor_map[new_out.name] = new_out node['outputs'][i] = new_out.name op.outputs[i] = new_out reshape_op = tfl.ReshapeOperator([new_out, shape_tensor], [op_out], shape_tensor.tensor) reshape_op.extra_hints['direction'] = 'down' self.graph.add_operator(reshape_op) tensor_node_dict[op_out.name] = self.graph.graph.vs.find(name=self.graph.tensor_node_map[op_out.name]) if node['node_type'] == ExtendedOperator.UNPACK: prev_shape = tmp_prev_shape # OP specific dim handling logic if node['node_type'] in ( ExtendedOperator.CONCATENATION, ExtendedOperator.GATHER, ExtendedOperator.UNPACK, ExtendedOperator.PACK, ): new_axis = prev_shape.index(-1) op.axis = new_axis elif node['node_type'] == ExtendedOperator.SPLIT_V: new_dim = prev_shape.index(-1) new_dim_tensor = self.create_attr_tensor(np.array(new_dim, dtype='int32')) actions.append((self.graph.replace_operator_input, (node, 2, new_dim_tensor, True))) elif node['node_type'] == ExtendedOperator.SPLIT: new_dim = prev_shape.index(-1) new_dim_tensor = self.create_attr_tensor(np.array(new_dim, dtype='int32')) actions.append((self.graph.replace_operator_input, (node, 0, new_dim_tensor, True))) elif node['node_type'] in (ExtendedOperator.PAD, ExtendedOperator.PADV2, ExtendedOperator.MIRROR_PAD): old_pad = op.inputs[1].tensor new_dim = prev_shape.index(-1) old_dim = next_shape.index(-1) new_pad = np.zeros((len(prev_shape), 2), dtype='int32') new_pad[new_dim, :] = old_pad[old_dim, :] new_pad_tensor = self.create_attr_tensor(new_pad) actions.append((self.graph.replace_operator_input, (node, 1, new_pad_tensor, True))) elif node['node_type'] == ExtendedOperator.PRELU: old_weight = op.inputs[1].tensor if old_weight.ndim != 1: new_dim = prev_shape.index(-1) old_dim = next_shape.index(-1) new_shape = np.ones(len(prev_shape) - 1, dtype='int32') new_shape[new_dim - 1] = old_weight.shape[old_dim - 1] new_shape_t = self.create_attr_tensor(new_shape) new_weight = self.create_transform_tensor(np.reshape(old_weight, new_shape)) self.graph.add_operator(tfl.ReshapeOperator([op.inputs[1], new_shape_t], [new_weight], new_shape)) actions.append((self.graph.replace_operator_input, (node, 1, new_weight, True))) elif node['node_type'] == ExtendedOperator.SLICE: new_dim = prev_shape.index(-1) old_dim = next_shape.index(-1) new_start = np.zeros(len(prev_shape), dtype='int32') new_start[new_dim] = op.inputs[1].tensor[old_dim] new_start_t = self.create_attr_tensor(new_start) new_size = np.array(prev_shape, dtype='int32') new_size[new_dim] = op.inputs[2].tensor[old_dim] new_size_t = self.create_attr_tensor(new_size) actions.append((self.graph.replace_operator_input, (node, 1, new_start_t, True))) actions.append((self.graph.replace_operator_input, (node, 2, new_size_t, True))) elif node['node_type'] == ExtendedOperator.STRIDED_SLICE: new_dim = prev_shape.index(-1) old_dim = next_shape.index(-1) new_start = np.zeros(len(prev_shape), dtype='int32') new_start[new_dim] = op.inputs[1].tensor[old_dim] if op.inputs[1].buffer is None: new_start_t = self.create_transform_tensor(new_start) starts_mid = new_start[new_dim : new_dim + 1] starts_mid_tensor = self.create_transform_tensor(starts_mid) slice_inputs = [ op.inputs[1], self.create_attr_tensor(np.array([old_dim], dtype='int32')), self.create_attr_tensor(np.array([1], dtype='int32')), ] self.graph.add_operator(tfl.SliceOperator(slice_inputs, [starts_mid_tensor])) starts_left = new_start[:new_dim] starts_right = new_start[new_dim + 1 :] starts_tensors = [] if len(starts_left) > 0: starts_tensors.append(self.create_attr_tensor(starts_left)) starts_tensors.append(starts_mid_tensor) if len(starts_right) > 0: starts_tensors.append(self.create_attr_tensor(starts_right)) if len(starts_tensors) > 1: self.graph.add_operator(tfl.ConcatenationOperator(starts_tensors, [new_start_t], 0)) else: new_start_t = starts_tensors[0] else: new_start_t = self.create_attr_tensor(new_start) new_end = np.array(prev_shape, dtype='int32') new_end[new_dim] = op.inputs[2].tensor[old_dim] if op.inputs[2].buffer is None: new_end_t = self.create_transform_tensor(new_end) ends_mid = new_end[new_dim : new_dim + 1] ends_mid_tensor = self.create_transform_tensor(ends_mid) slice_inputs = [ op.inputs[2], self.create_attr_tensor(np.array([old_dim], dtype='int32')), self.create_attr_tensor(np.array([1], dtype='int32')), ] self.graph.add_operator(tfl.SliceOperator(slice_inputs, [ends_mid_tensor])) ends_left = new_end[:new_dim] ends_right = new_end[new_dim + 1 :] ends_tensors = [] if len(ends_left) > 0: ends_tensors.append(self.create_attr_tensor(ends_left)) ends_tensors.append(ends_mid_tensor) if len(ends_right) > 0: ends_tensors.append(self.create_attr_tensor(ends_right)) if len(ends_tensors) > 1: self.graph.add_operator(tfl.ConcatenationOperator(ends_tensors, [new_end_t], 0)) else: new_end_t = ends_tensors[0] else: new_end_t = self.create_attr_tensor(new_end) new_stride = np.ones(len(prev_shape), dtype='int32') new_stride[new_dim] = op.inputs[3].tensor[old_dim] new_stride_t = self.create_attr_tensor(new_stride) actions.append((self.graph.replace_operator_input, (node, 1, new_start_t, True))) actions.append((self.graph.replace_operator_input, (node, 2, new_end_t, True))) actions.append((self.graph.replace_operator_input, (node, 3, new_stride_t, True))) elif node['node_type'] == ExtendedOperator.TILE: old_shape = op.inputs[1].tensor new_dim = prev_shape.index(-1) old_dim = next_shape.index(-1) new_shape = np.ones(len(prev_shape), dtype='int32') new_shape[new_dim] = old_shape[old_dim] new_shape_tensor = self.create_attr_tensor(new_shape) actions.append((self.graph.replace_operator_input, (node, 1, new_shape_tensor, True))) elif node['node_type'] in ( ExtendedOperator.SUM, ExtendedOperator.ARG_MIN, ExtendedOperator.ARG_MAX, ExtendedOperator.REDUCE_MIN, ExtendedOperator.REDUCE_MAX, ExtendedOperator.REDUCE_PROD, ExtendedOperator.MEAN, ): new_axis = prev_shape.index(-1) axis_arr = np.array([new_axis], dtype='int32') axis_tensor = self.create_attr_tensor(axis_arr) actions.append((self.graph.replace_operator_input, (node, 1, axis_tensor, True))) elif dim_indice is not None: raise NotImplementedError(f'{node["node_type"]} has the property `dims` but is not handled') for edge in next_edges: source = tensor_node_dict[edge['name']] self.graph.graph.add_edge(source, edge.target_vertex, name=edge['name'], label=edge['name']) # Process actions ids = [] for func, args in actions: node = args[0] res = func(*args) if res is not None: ids.extend(res) remove_edges = list(set(remove_edges + ids)) self.graph.graph.delete_edges(remove_edges) self.graph.graph.delete_vertices(remove_vertices) return num_actions @class_conditional(lambda self: self.level >= GraphOptimizer.COMMON_OPTIMIZE) def fuse_bmm_add_pass(self): edges = self.graph.graph.es.select(functools.partial(is_bmm_add_edge, graph_converter=self.graph.graph)) filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in edges] filtered_pairs = [ p for p in filtered_pairs if p[0]['node_type'] != ExtendedOperator.FULLY_CONNECTED or len(p[0]['op'].inputs) == 2 or not np.any(p[0]['op'].inputs[2].tensor) ] remove_ids = [] ops = [] restore_mapping = [] for bmm, add in filtered_pairs: restore_nodes = [] # For each node that is next of a transformable node, # a. if it is an output node, remove it anyway since it will always be reconstructed # b. otherwise, record the info of the edge so that we may restore it after reconstruction for out_edge in add.out_edges(): next_node = self.graph.graph.vs[out_edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: remove_ids.append(next_node.index) del self.graph.tensor_map[next_node['outputs'][0]] del self.graph.tensor_node_map[next_node['outputs'][0]] else: restore_nodes.append((out_edge['name'], next_node['name'])) # Remove the mapping since they are going to be removed for output_name in add['outputs']: del self.graph.tensor_map[output_name] del self.graph.tensor_node_map[output_name] restore_mapping.append(restore_nodes) ops.append((bmm, add)) remove_ids.append(bmm.index) remove_ids.append(add.index) # Make sure the nodes are topologically sorted sorted_ops = [ (nodes[0]['op'], nodes[1]['op']) for nodes in sorted(ops, key=lambda x: int(re.search(r'\d+', x[1]['name'])[0])) ] # Delete nodes before transformation in the graph self.graph.graph.delete_vertices(remove_ids) for (bmm, add), mapping in zip(sorted_ops, restore_mapping): input_tensor = bmm.inputs[0] weight_tensor = bmm.inputs[1] bias_tensor = add.inputs[1] output_tensor = add.outputs[0] ops = [] if isinstance(bmm, tfl.BatchMatmulOperator): weight_t = self.create_transform_tensor(np.transpose(weight_tensor.tensor)) weight_perm = self.create_attr_tensor(np.array([1, 0], dtype='int32')) ops.append(tfl.TransposeOperator([weight_tensor, weight_perm], [weight_t])) else: weight_t = weight_tensor keep_dims = output_tensor.tensor.ndim > 2 ops.append( tfl.FullyConnectedOperator( [input_tensor, weight_t, bias_tensor], [output_tensor], fusedActivationFunction=add.fusedActivationFunction, keepNumDims=keep_dims, ) ) for op in ops: self.graph.add_operator(op, transform=True) self.graph.try_restore_edges(mapping) @class_conditional(lambda self: self.max_transpose_dims > 0) def lower_transpose_dim_pass(self): vertices = self.graph.graph.vs.select( functools.partial( is_high_dim_transpose_node, graph_converter=self.graph.graph, max_transpose_dims=self.max_transpose_dims ) ) remove_ids = [] ops = [] restore_mapping = [] for trans in vertices: restore_nodes = [] # For each node that is next of a transformable node, # a. if it is an output node, remove it anyway since it will always be reconstructed # b. otherwise, record the info of the edge so that we may restore it after reconstruction for out_edge in trans.out_edges(): next_node = self.graph.graph.vs[out_edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: remove_ids.append(next_node.index) del self.graph.tensor_map[next_node['outputs'][0]] del self.graph.tensor_node_map[next_node['outputs'][0]] else: restore_nodes.append((out_edge['name'], next_node['name'])) # Remove the mapping since they are going to be removed for output_name in trans['outputs']: del self.graph.tensor_map[output_name] del self.graph.tensor_node_map[output_name] restore_mapping.append(restore_nodes) remove_ids.append(trans.index) # Make sure the nodes are topologically sorted sorted_ops = [node['op'] for node in sorted(vertices, key=lambda x: int(re.search(r'\d+', x['name'])[0]))] # Delete nodes before transformation in the graph self.graph.graph.delete_vertices(remove_ids) for trans, mapping in zip(sorted_ops, restore_mapping): input_tensor = trans.inputs[0] perm_tensor = trans.inputs[1] output_tensor = trans.outputs[0] input_shape = input_tensor.shape perm = perm_tensor.tensor output_shape = output_tensor.shape last_perm = None last_dim = None cum_dim = None new_shape = [] new_perm = [] for d, p in zip(input_shape, perm): if last_dim is None and last_perm is None: cum_dim = d else: if p - last_perm == 1 or d == 1 or cum_dim == 1: cum_dim *= d else: new_shape.append(cum_dim) new_perm.append(last_perm) cum_dim = d last_dim = d last_perm = p new_shape.append(cum_dim) new_perm.append(last_perm) new_perm_arr = np.argsort(new_perm).astype('int32') assert ( len(new_shape) <= self.max_transpose_dims ), f"Don't know how to reduce the number of dims of transpose with input shape {input_shape}, perm {perm}" ops = [] input_reduced = self.create_transform_tensor( np.reshape(input_tensor.tensor, new_shape), quantization=input_tensor.quantization ) reduced_shape = self.create_attr_tensor(np.array(new_shape, dtype='int32')) ops.append(tfl.ReshapeOperator([input_tensor, reduced_shape], [input_reduced], new_shape)) transposed = self.create_transform_tensor( np.transpose(input_reduced.tensor, new_perm_arr), quantization=input_tensor.quantization ) new_perm_tensor = self.create_attr_tensor(np.array(new_perm_arr, dtype='int32')) ops.append(tfl.TransposeOperator([input_reduced, new_perm_tensor], [transposed])) output_shape_tensor = self.create_attr_tensor(np.array(output_shape, dtype='int32')) ops.append(tfl.ReshapeOperator([transposed, output_shape_tensor], [output_tensor], output_shape)) for op in ops: self.graph.add_operator(op, transform=True) self.graph.try_restore_edges(mapping) @class_conditional(lambda self: self.group_conv_rewrite) def group_conv_rewrite_pass(self): vertices = self.graph.graph.vs.select(functools.partial(is_group_conv_node, graph_converter=self.graph.graph)) remove_ids = [] ops = [] restore_mapping = [] for conv in vertices: restore_nodes = [] # For each node that is next of a transformable node, # a. if it is an output node, remove it anyway since it will always be reconstructed # b. otherwise, record the info of the edge so that we may restore it after reconstruction for out_edge in conv.out_edges(): next_node = self.graph.graph.vs[out_edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: remove_ids.append(next_node.index) del self.graph.tensor_map[next_node['outputs'][0]] del self.graph.tensor_node_map[next_node['outputs'][0]] else: restore_nodes.append((out_edge['name'], next_node['name'])) # Remove the mapping since they are going to be removed for output_name in conv['outputs']: del self.graph.tensor_map[output_name] del self.graph.tensor_node_map[output_name] restore_mapping.append(restore_nodes) remove_ids.append(conv.index) # Make sure the nodes are topologically sorted sorted_ops = [node['op'] for node in sorted(vertices, key=lambda x: int(re.search(r'\d+', x['name'])[0]))] # Delete nodes before transformation in the graph self.graph.graph.delete_vertices(remove_ids) for conv, mapping in zip(sorted_ops, restore_mapping): input_tensor = conv.inputs[0] weight_tensor = conv.inputs[1] bias_tensor = conv.inputs[2] if len(conv.inputs) > 2 else None output_tensor = conv.outputs[0] num_input_channel = input_tensor.shape[3] num_weight_channel = weight_tensor.shape[3] num_chunks = num_input_channel // num_weight_channel ops = [] input_tensors = [ self.create_transform_tensor(arr, quantization=input_tensor.quantization) for arr in np.split(input_tensor.tensor, num_chunks, 3) ] output_tensors = [ self.create_transform_tensor(arr, quantization=output_tensor.quantization) for arr in np.split(output_tensor.tensor, num_chunks, 3) ] weights = [ self.create_attr_tensor(arr, quantization=weight_tensor.quantization) for arr in np.split(weight_tensor.tensor, num_chunks, 0) ] if bias_tensor is not None: biases = [ self.create_attr_tensor(arr, quantization=bias_tensor.quantization) for arr in np.split(bias_tensor.tensor, num_chunks, 0) ] else: biases = [None] * num_chunks dim_tensor = self.create_attr_tensor(np.array(3, dtype='int32')) ops.append(tfl.SplitOperator([dim_tensor, input_tensor], input_tensors, num_chunks)) for it, ot, w, b in zip(input_tensors, output_tensors, weights, biases): inputs = [it, w] if b is not None: inputs.append(b) ops.append( tfl.Conv2dOperator( inputs, [ot], strideH=conv.strideH, strideW=conv.strideW, dilationHFactor=conv.dilationHFactor, dilationWFactor=conv.dilationWFactor, fusedActivationFunction=conv.fusedActivationFunction, padding=conv.padding, ) ) ops.append(tfl.ConcatenationOperator(output_tensors, [output_tensor], 3)) for op in ops: self.graph.add_operator(op, transform=True) self.graph.try_restore_edges(mapping) @class_conditional(lambda self: self.group_conv_rewrite) def group_deconv_rewrite_pass(self): vertices = self.graph.graph.vs.select(functools.partial(is_group_deconv_node, graph_converter=self.graph.graph)) remove_ids = [] ops = [] restore_mapping = [] for conv in vertices: restore_nodes = [] # For each node that is next of a transformable node, # a. if it is an output node, remove it anyway since it will always be reconstructed # b. otherwise, record the info of the edge so that we may restore it after reconstruction for out_edge in conv.out_edges(): next_node = self.graph.graph.vs[out_edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: remove_ids.append(next_node.index) del self.graph.tensor_map[next_node['outputs'][0]] del self.graph.tensor_node_map[next_node['outputs'][0]] else: restore_nodes.append((out_edge['name'], next_node['name'])) # Remove the mapping since they are going to be removed for output_name in conv['outputs']: del self.graph.tensor_map[output_name] del self.graph.tensor_node_map[output_name] restore_mapping.append(restore_nodes) remove_ids.append(conv.index) # Make sure the nodes are topologically sorted sorted_ops = [node['op'] for node in sorted(vertices, key=lambda x: int(re.search(r'\d+', x['name'])[0]))] # Delete nodes before transformation in the graph self.graph.graph.delete_vertices(remove_ids) for conv, mapping in zip(sorted_ops, restore_mapping): input_tensor = conv.inputs[2] weight_tensor = conv.inputs[1] output_shape_tensor = conv.inputs[0] bias_tensor = conv.inputs[3] if len(conv.inputs) > 3 else None output_tensor = conv.outputs[0] num_output_channel = output_tensor.shape[3] num_weight_channel = weight_tensor.shape[0] num_chunks = num_output_channel // num_weight_channel ops = [] input_tensors = [ self.create_transform_tensor(arr, quantization=input_tensor.quantization) for arr in np.split(input_tensor.tensor, num_chunks, 3) ] output_tensors = [ self.create_transform_tensor(arr, quantization=output_tensor.quantization) for arr in np.split(output_tensor.tensor, num_chunks, 3) ] weights = [ self.create_attr_tensor(arr, quantization=weight_tensor.quantization) for arr in np.split(weight_tensor.tensor, num_chunks, 3) ] if bias_tensor is not None: biases = [ self.create_attr_tensor(arr, quantization=bias_tensor.quantization) for arr in np.split(bias_tensor.tensor, num_chunks, 0) ] else: biases = [None] * num_chunks new_os = output_shape_tensor.tensor.copy() new_os[3] = num_weight_channel new_ost = self.create_attr_tensor(new_os) dim_tensor = self.create_attr_tensor(np.array(3, dtype='int32')) ops.append(tfl.SplitOperator([dim_tensor, input_tensor], input_tensors, num_chunks)) for it, ot, w, b in zip(input_tensors, output_tensors, weights, biases): inputs = [new_ost, w, it] if b is not None: inputs.append(b) ops.append( tfl.TransposeConvOperator( inputs, [ot], padding=conv.padding, strideH=conv.strideH, strideW=conv.strideW, ) ) ops.append(tfl.ConcatenationOperator(output_tensors, [output_tensor], 3)) for op in ops: self.graph.add_operator(op, transform=True) self.graph.try_restore_edges(mapping) @class_conditional(lambda self: self.tflite_micro_rewrite) def cat_split_pass(self): vertices = self.graph.graph.vs.select(functools.partial(is_large_cat_node, graph_converter=self.graph.graph)) remove_ids = [] ops = [] restore_mapping = [] for cat in vertices: restore_nodes = [] # For each node that is next of a transformable node, # a. if it is an output node, remove it anyway since it will always be reconstructed # b. otherwise, record the info of the edge so that we may restore it after reconstruction for out_edge in cat.out_edges(): next_node = self.graph.graph.vs[out_edge.target] if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: remove_ids.append(next_node.index) del self.graph.tensor_map[next_node['outputs'][0]] del self.graph.tensor_node_map[next_node['outputs'][0]] else: restore_nodes.append((out_edge['name'], next_node['name'])) # Remove the mapping since they are going to be removed for output_name in cat['outputs']: del self.graph.tensor_map[output_name] del self.graph.tensor_node_map[output_name] restore_mapping.append(restore_nodes) remove_ids.append(cat.index) # Make sure the nodes are topologically sorted sorted_ops = [node['op'] for node in sorted(vertices, key=lambda x: int(re.search(r'\d+', x['name'])[0]))] # Delete nodes before transformation in the graph self.graph.graph.delete_vertices(remove_ids) for cat, mapping in zip(sorted_ops, restore_mapping): input_tensors = cat.inputs layer_inputs = input_tensors output_tensor = cat.outputs[0] axis = cat.axis last_layer = False ops = [] while True: layer_outputs = [] while len(layer_inputs) > 0: curr_inputs = layer_inputs[:10] input_arrs = [t.tensor for t in curr_inputs] output_arr = np.concatenate(input_arrs, axis) if last_layer: curr_output = output_tensor else: curr_output = self.create_transform_tensor(output_arr, quantization=output_tensor.quantization) layer_outputs.append(curr_output) ops.append(tfl.ConcatenationOperator(curr_inputs, [curr_output], axis)) layer_inputs = layer_inputs[10:] if len(layer_outputs) == 0: break elif len(layer_outputs) <= 10: last_layer = True layer_inputs = layer_outputs for op in ops: self.graph.add_operator(op, transform=True) self.graph.try_restore_edges(mapping) def input_transpose_pass(self): nhwc2nchw_perm = np.array([0, 3, 1, 2], dtype='int32') nchw2nhwc_perm = np.array([0, 2, 3, 1], dtype='int32') remove_edges = [] for name, transpose in zip(self.graph.inputs, self.graph.input_transpose): if transpose is True: node_name = self.graph.tensor_node_map[name] node = self.graph.graph.vs.find(name=node_name) assert node['node_type'] == ExtendedOperator.INPUT_NODE # For quantized graphs, we insert the transpose op after the quantize op next_node = None if node.outdegree() == 1: next_node = node.out_edges()[0].target_vertex if next_node['node_type'] != ExtendedOperator.QUANTIZE: next_node = None # Transpose input tensor shapes input_tensor = self.graph.tensor_map[node['name']] input_tensor.tensor = np.transpose(input_tensor.tensor, nchw2nhwc_perm) input_tensor.shape = input_tensor.tensor.shape # Transpose quantize output tensor shapes last_tensor = input_tensor last_node = node if next_node is not None: last_node = next_node last_tensor = next_node['op'].outputs[0] last_tensor.tensor = np.transpose(last_tensor.tensor, nchw2nhwc_perm) last_tensor.shape = last_tensor.tensor.shape # Create new transpose op nhwc2nchw_perm_tensor = self.create_attr_tensor(nhwc2nchw_perm) transposed = self.create_transform_tensor( np.transpose(last_tensor.tensor, nhwc2nchw_perm), quantization=last_tensor.quantization ) transpose_op = tfl.TransposeOperator([last_tensor, nhwc2nchw_perm_tensor], [transposed]) transpose_op.extra_hints['direction'] = 'down' self.graph.add_operator(transpose_op) # Get the newly-generated node new_node_name = self.graph.tensor_node_map[transposed.name] new_node = self.graph.graph.vs.find(name=new_node_name) # Connect the transpose op to the graph self.graph.replace_next_tensors(last_node, new_node, transposed.name, [new_node_name]) # Collect the unused connections for edge in last_node.out_edges(): target_vertex = edge.target_vertex if target_vertex['name'] != new_node_name: remove_edges.append(edge.index) # Remove the collected edges self.graph.graph.delete_edges(remove_edges) @class_conditional(lambda self: self.quantize_input_output_type is not None) def quantize_input_output_type_pass(self): remove_edges = [] remove_vertices = [] for i, name in enumerate(self.graph.inputs): if self.fuse_input_indices is not None: if i not in self.fuse_input_indices: continue node_name = self.graph.tensor_node_map[name] node = self.graph.graph.vs.find(name=node_name) assert node['node_type'] == ExtendedOperator.INPUT_NODE # Update input tensor input_tensor = self.graph.tensor_map[node['outputs'][0]] input_type = str(input_tensor.dtype) if input_type == self.quantize_input_output_type: continue input_arr = input_tensor.tensor.copy() input_quantization = copy.deepcopy(input_tensor.quantization) if input_type == 'int8' and self.quantize_input_output_type == 'uint8': input_tensor.tensor = (input_tensor.tensor.astype('int32') + 128).astype('uint8') input_tensor.quantization.zero_point += 128 input_tensor.dtype = input_tensor.tensor.dtype elif input_type == 'uint8' and self.quantize_input_output_type == 'int8': input_tensor.tensor = (input_tensor.tensor.astype('int32') - 128).astype('int8') input_tensor.quantization.zero_point -= 128 input_tensor.dtype = input_tensor.tensor.dtype else: raise AssertionError( f'Unsupported types: input_type: {input_type}, quantize_input_type:' f' {self.quantize_input_output_type}' ) # Create new quantize op requantized = self.create_transform_tensor(input_arr, quantization=input_quantization) quantize_op = tfl.QuantizeOperator([input_tensor], [requantized]) self.graph.add_operator(quantize_op) # Get the newly-generated node new_node_name = self.graph.tensor_node_map[requantized.name] new_node = self.graph.graph.vs.find(name=new_node_name) # Connect the quantize op to the graph self.graph.replace_next_tensors(node, new_node, requantized.name, [new_node_name]) # Collect the unused connections for edge in node.out_edges(): target_vertex = edge.target_vertex if target_vertex['name'] != new_node_name: remove_edges.append(edge.index) output_mapping = {} for i, name in enumerate(self.graph.outputs): if self.fuse_output_indices is not None: if i not in self.fuse_output_indices: continue output_tensor = self.graph.tensor_map[name] output_type = str(output_tensor.dtype) if output_type == self.quantize_input_output_type: continue node_name = self.graph.tensor_node_map[name] node = self.graph.graph.vs.find(name=node_name) for edge in node.out_edges(): next_node = edge.target_vertex if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: remove_vertices.append(next_node.index) # Update output tensor output_arr = output_tensor.tensor.copy() output_quantization = copy.deepcopy(output_tensor.quantization) if output_type == 'int8' and self.quantize_input_output_type == 'uint8': output_arr = (output_arr.astype('int32') + 128).astype('uint8') output_quantization.zero_point += 128 elif output_type == 'uint8' and self.quantize_input_output_type == 'int8': output_arr = (output_arr.astype('int32') - 128).astype('int8') output_quantization.zero_point -= 128 else: raise AssertionError( f'Unsupported types: output_type: {output_type}, quantize_input_type:' f' {self.quantize_input_output_type}' ) requantized = self.create_transform_tensor(output_arr, quantization=output_quantization) quantize_op = tfl.QuantizeOperator([output_tensor], [requantized]) self.graph.add_operator(quantize_op) output_mapping[name] = requantized.name if len(output_mapping) > 0: new_outputs = [] output_names = [] for name in self.graph.outputs: if name in output_mapping: new_outputs.append(output_mapping[name]) output_names.append(output_mapping[name]) else: new_outputs.append(name) self.graph.outputs.clear() self.graph.outputs.extend(new_outputs) self.graph.add_outputs(output_names) # Remove the collected edges & vertices self.graph.graph.delete_edges(remove_edges) self.graph.graph.delete_vertices(remove_vertices) def output_transpose_pass(self): nhwc2nchw_perm = np.array([0, 3, 1, 2], dtype='int32') nchw2nhwc_perm = np.array([0, 2, 3, 1], dtype='int32') if isinstance(self.graph.output_transpose, (list, tuple)): assert len(self.graph.output_transpose) == len(self.graph.outputs) else: self.graph.output_transpose = [self.graph.output_transpose] * len(self.graph.outputs) filtered_dict = {} for i, (name, transpose) in enumerate(zip(self.graph.outputs, self.graph.output_transpose)): if name in filtered_dict: old_transpose = filtered_dict[name] assert ( transpose == old_transpose ), f"outputs {i} points to an exising tensor {name}, but their property `output_transpose` is different" else: filtered_dict[name] = transpose prev_modify_node_indices = {} prev_modify_next_indices = {} next_modify_node_indices = {} for name, transpose in filtered_dict.items(): if name in self.graph.tensor_map: tensor = self.graph.tensor_map[name] if transpose is None: transpose = len(tensor.shape) == 4 else: transpose = False for i, n in enumerate(self.graph.outputs): if name == n: self.graph.output_transpose[i] = transpose if transpose: node_name = self.graph.tensor_node_map[name] node = self.graph.graph.vs.find(name=node_name) tensor_idx = node['outputs'].index(name) prev_node = None if node['node_type'] == ExtendedOperator.DEQUANTIZE: prev_node_name = self.graph.tensor_node_map[node['op'].inputs[0].name] prev_node = self.graph.graph.vs.find(name=prev_node_name) if prev_node is None: next_modify_node_indices.setdefault(node, set()) next_modify_node_indices[node].add(tensor_idx) else: prev_modify_node_indices.setdefault(node, set()) prev_modify_node_indices[node].add(0) prev_modify_next_indices.setdefault(node, set()) prev_modify_next_indices[node].add(tensor_idx) remove_edges = [] remove_vertices = [] actions = [] for node, index in prev_modify_node_indices.items(): next_indices = prev_modify_next_indices[node] op = node['op'] tensor_names = [node['outputs'][i] for i in index] next_nodes = {} for edge in node.out_edges(): if edge['label'] not in tensor_names: continue if edge.index in remove_edges: continue tensor_idx = tensor_names.index(edge['label']) next_node = self.graph.graph.vs[edge.target] if next_node['node_type'] not in (ExtendedOperator.OUTPUT_NODE, ExtendedOperator.UNUSED_NODE): next_nodes.setdefault(tensor_idx, []) next_nodes[tensor_idx].append(next_node) prev_nodes = [] prev_output_indices = [] for i in index: prev_node_name = op.inputs[i].name prev_node = self.graph.graph.vs.find(name=self.graph.tensor_node_map[prev_node_name]) prev_nodes.append(prev_node) prev_output_indices.append(prev_node['outputs'].index(prev_node_name)) tensor_node_dict = {} for prev_node, prev_idx, next_idx in zip(prev_nodes, index, prev_output_indices): if prev_node['op'] is None: prev_out = self.graph.tensor_map[prev_node['outputs'][0]] else: prev_out = prev_node['op'].outputs[next_idx] if prev_out.name in tensor_node_dict: prev_new_out, skip = tensor_node_dict[prev_out.name] actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True, skip))) skip += 1 tensor_node_dict[prev_out.name] = (prev_new_out, skip) else: perm_tensor = self.create_attr_tensor(nchw2nhwc_perm) prev_new_out = self.create_transform_tensor( np.transpose(prev_out.tensor, nchw2nhwc_perm), quantization=prev_out.quantization ) tensor_node_dict[prev_out.name] = (prev_new_out, 1) prev_transpose_op = tfl.TransposeOperator([prev_out, perm_tensor], [prev_new_out]) prev_transpose_op.extra_hints['direction'] = 'up' self.graph.add_operator(prev_transpose_op) actions.append((self.graph.replace_operator_input, (node, prev_idx, prev_new_out, True))) tensor_mapping = {} for i in next_indices: t = op.outputs[i] t.tensor = np.transpose(t.tensor, nchw2nhwc_perm) t.shape = t.tensor.shape if i in next_nodes: new_t = self.create_transform_tensor(np.transpose(t.tensor, nhwc2nchw_perm)) perm_t = self.create_attr_tensor(nhwc2nchw_perm) next_transpose_op = tfl.TransposeOperator([t, perm_t], [new_t]) next_transpose_op.extra_hints['direction'] = 'down' self.graph.add_operator(next_transpose_op) tensor_mapping[t.name] = new_t for nodes in next_nodes.values(): for n in nodes: next_op = n['op'] for i, t in enumerate(next_op.inputs): if t.name in tensor_mapping: actions.append((self.graph.replace_operator_input, (n, i, tensor_mapping[t.name]))) for node, index in next_modify_node_indices.items(): op = node['op'] tensor_names = [node['outputs'][i] for i in index] out_nodes = [] next_nodes = [] next_edges = [] for edge in node.out_edges(): if edge['label'] not in tensor_names: continue if edge.index in remove_edges: continue next_node = self.graph.graph.vs[edge.target] tensor_idx = tensor_names.index(edge['label']) if next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: out_nodes.append(next_node) elif next_node['node_type'] != ExtendedOperator.UNUSED_NODE: next_nodes.append(next_node) next_edges.append(edge) remove_vertices.extend([x.index for x in out_nodes]) remove_edges.extend([x.index for x in next_edges]) for n in out_nodes: del self.graph.tensor_map[n['outputs'][0]] del self.graph.tensor_node_map[n['outputs'][0]] tensor_node_dict = {} for i, op_out in enumerate(op.outputs): if i not in index: continue op_out.tensor = np.transpose(op_out.tensor, nchw2nhwc_perm) op_out.shape = op_out.tensor.shape perm_tensor = self.create_attr_tensor(nchw2nhwc_perm) new_out = self.create_transform_tensor( np.transpose(op_out.tensor, nhwc2nchw_perm), quantization=op_out.quantization ) # Update relations if op_out.name in self.graph.tensor_node_map: del self.graph.tensor_node_map[op_out.name] self.graph.tensor_node_map[new_out.name] = node['name'] self.graph.tensor_map[new_out.name] = new_out node['outputs'][i] = new_out.name op.outputs[i] = new_out next_transpose_op = tfl.TransposeOperator([new_out, perm_tensor], [op_out]) next_transpose_op.extra_hints['direction'] = 'up' self.graph.add_operator(next_transpose_op) tensor_node_dict[op_out.name] = ( self.graph.graph.vs.find(name=self.graph.tensor_node_map[new_out.name]), new_out.name, ) # Connect next edges and replace next tensors for edge in next_edges: old_name = edge['name'] source, new_name = tensor_node_dict[old_name] target = edge.target_vertex self.graph.graph.add_edge(source, target, name=new_name, label=new_name) op = target['op'] for i, op_input in enumerate(op.inputs): if op_input.name == old_name: op.inputs[i] = self.graph.tensor_map[new_name] break # Process actions ids = [] for func, args in actions: node = args[0] res = func(*args) if res is not None: ids.extend(res) remove_edges = list(set(remove_edges + ids)) self.graph.graph.delete_edges(remove_edges) self.graph.graph.delete_vertices(remove_vertices) def connect_unused_tensors_pass(self): filtered_nodes = self.graph.graph.vs.select( functools.partial(is_multi_output_op_node, graph_converter=self.graph.graph) ) list_unpack_names = set([i for s in self.graph.iterable_map.values() for i in s]) all_tensors = set(self.graph.graph.es['label']) names = [] for node in filtered_nodes: output_names = node['outputs'] # Recognizes the pattern SPLIT -> (RESHAPE, ..., RESHAPE) if not list_unpack_names.isdisjoint(set(output_names)): output_names = [] outdegree = 0 for edge in node.out_edges(): target_vertex = edge.target_vertex if target_vertex['node_type'] == ExtendedOperator.RESHAPE: outdegree += target_vertex.outdegree() output_names.append(target_vertex['outputs'][0]) # Only nodes with partially unused tensors are supported if outdegree == 0: continue for out in output_names: if out not in all_tensors: names.append(out) self.graph.add_outputs(names, ExtendedOperator.UNUSED_NODE) def output_list_unpack_pass(self): output_names = [] unpacked_outputs = [] for name in self.graph.outputs: if name in self.graph.iterable_map: names = self.graph.get_list_expanded_names(name) unpacked_outputs.extend(names) output_names.extend(names) else: unpacked_outputs.append(name) self.graph.outputs.clear() self.graph.outputs.extend(unpacked_outputs) self.graph.add_outputs(output_names) @class_conditional(lambda self: self.fuse_quant) def fuse_quant_dequant_nodes(self): edges = self.graph.graph.es.select(functools.partial(is_quant_dequant_edge, graph_converter=self.graph.graph)) filtered_pairs = [[self.graph.graph.vs[x.source], self.graph.graph.vs[x.target]] for x in edges] remove_vertices = [] input_mapping = {} output_mapping = {} for prev, next in filtered_pairs: if prev['node_type'] == ExtendedOperator.INPUT_NODE: input_name = prev['outputs'][0] if self.fuse_input_indices is not None: input_idx = self.graph.inputs.index(input_name) if input_idx not in self.fuse_input_indices: continue remove_vertices.append(prev) next['node_type'] = prev['node_type'] next['op'] = None input_mapping.setdefault(input_name, []) input_mapping[input_name].extend(next['outputs']) else: if prev['op'] is not None: prev_name = prev['op'].outputs[0].name if self.fuse_output_indices is not None: output_idx = self.graph.outputs.index(prev_name) if output_idx not in self.fuse_output_indices: continue prev['node_type'] = next['node_type'] new_name = prev['op'].inputs[0].name prev['op'] = None output_mapping.setdefault(prev_name, []) output_mapping[prev_name].append(new_name) remove_vertices.append(next) self.graph.graph.delete_vertices(remove_vertices) if len(input_mapping) > 0: new_inputs = [] for name in self.graph.inputs: if name in input_mapping: new_inputs.extend(input_mapping[name]) else: new_inputs.append(name) self.graph.inputs.clear() self.graph.inputs.extend(new_inputs) if len(output_mapping) > 0: new_outputs = [] for name in self.graph.outputs: if name in output_mapping: new_outputs.extend(output_mapping[name]) else: new_outputs.append(name) self.graph.outputs.clear() self.graph.outputs.extend(new_outputs) def optimize(self): # Input/output passes self.output_list_unpack_pass() self.input_transpose_pass() self.output_transpose_pass() # Connect unused tensors with special nodes self.connect_unused_tensors_pass() # Transpose, Reshape and NO-OP cleanup self.branch_reshape_expand_pass() self.fuse_simple_reshape_pass() self.branch_transpose_expand_pass() self.fuse_simple_transpose_pass() self.fuse_simple_gather_pass() for branch in (False, True): self.remove_noop_pass(branch) self.fuse_wrapped_reshape_within_transpose_pass() # Buffer folding, which is needed by the fusion passes below for _ in range(2): self.fold_reshape_buffer() self.fold_transpose_buffer() # Move `transpose` ops for the rewrite quantizable pass self.elementwise_op_transpose_passthrough_pass(quantizable_ops_only=True) self.branch_transpose_expand_pass() self.fuse_simple_transpose_pass() # Fuse reciprocal and sqrt self.fuse_reciprocal_sqrt() # Map quantizable ops to quantized kernels self.elementwise_op_quantize_passthrough_pass() # Remove consecutive dequantize and quantize nodes self.fuse_dequant_quant_pass(q_first=False) # OP fusion passes before transformation self.fuse_conv_fc_bn() self.fuse_activation() self.fuse_requantize() self.fuse_bn_conv() # Convert TinyNeuralNetwork ops to TFLite ops self.transform_graph() # OP fusion passes after transformation self.fuse_bmm_add_pass() self.fuse_activation() # Transpose and reshape cleanup self.branch_reshape_expand_pass() self.branch_transpose_expand_pass() self.fuse_simple_transpose_pass() self.fuse_simple_reshape_pass() # Branch transpose & reshape cleanup for i in range(11): t_count = self.elementwise_op_transpose_passthrough_pass() self.branch_transpose_expand_pass() self.fuse_simple_transpose_pass() r_count = self.elementwise_op_reshape_passthrough_pass() self.branch_reshape_expand_pass() self.fuse_simple_reshape_pass() c_count = self.elementwise_reshape_transpose_passthrough_pass() self.branch_transpose_expand_pass() self.fuse_simple_transpose_pass() if t_count + r_count + c_count == 0: log.debug(f'elem p/t pass finished in {i + 1} steps') break # Other cleanups self.fuse_simple_slice_pass() for branch in (False, True): self.remove_noop_pass(branch) self.fuse_wrapped_reshape_within_transpose_pass() # Buffer folding for _ in range(2): self.fold_reshape_buffer() self.fold_transpose_buffer() # Transpose and reshape cleanup for _ in range(2): self.transpose_to_reshape_pass() self.branch_reshape_expand_pass() self.fuse_simple_reshape_pass() self.fuse_simple_transpose_pass() self.lower_transpose_dim_pass() # Some advanced fusion logic self.fuse_conv2d_gather() # Remove consecutive dequantize and quantize nodes self.fuse_dequant_quant_pass(q_first=True) # Fuse reciprocal and sqrt self.fuse_reciprocal_sqrt() # Remove additional tile nodes before elementwise ops self.remove_tile_before_binary_elementwise_ops() # Fuse activation self.fuse_activation() # Fuse quant/dequant nodes self.fuse_quant_dequant_nodes() # Input output quantize type self.quantize_input_output_type_pass() # Fuse same padding self.fuse_same_padding() self.fuse_same_padding_slicing() self.fuse_gather_conv2d() # Group conv & deconv self.group_conv_rewrite_pass() self.group_deconv_rewrite_pass() # TFLite micro specific self.cat_split_pass() self.split_requantize() # Group the same tensors into one self.group_tensors_pass() # Final cleanup self.cleanup_dead_nodes() def is_bn_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] in (ExtendedOperator.GENERIC_CONV, ExtendedOperator.GENERIC_DECONV, ExtendedOperator.FULLY_CONNECTED) and target_vertex['node_type'] == ExtendedOperator.BATCH_NORM and source_vertex.outdegree() == 1 and target_vertex['op'].inputs[1].buffer is not None and target_vertex['op'].inputs[2].buffer is not None and source_vertex['op'].inputs[1].buffer is not None and ( target_vertex['op'].fusedActivationFunction == ActivationFunctionType.NONE or source_vertex['op'].fusedActivationFunction in (ActivationFunctionType.NONE, target_vertex['op'].fusedActivationFunction) ) ) def is_rev_bn_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( target_vertex['node_type'] == ExtendedOperator.GENERIC_CONV and source_vertex['node_type'] == ExtendedOperator.BATCH_NORM and source_vertex.outdegree() == 1 and source_vertex['op'].inputs[1].buffer is not None and source_vertex['op'].inputs[2].buffer is not None and target_vertex['op'].inputs[1].buffer is not None and source_vertex['op'].fusedActivationFunction == ActivationFunctionType.NONE ) def is_padding_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] in (ExtendedOperator.PAD, ExtendedOperator.PADV2) and ( len(source_vertex['op'].inputs) == 2 or ( len(source_vertex['op'].inputs) == 3 and source_vertex['op'].inputs[2].dtype == np.dtype('float32') and ( ( source_vertex['op'].inputs[2].tensor[0] == 0.0 and target_vertex['node_type'] != ExtendedOperator.MAX_POOL_2D ) or ( source_vertex['op'].inputs[2].tensor[0] == np.finfo(np.float32).min and target_vertex['node_type'] == ExtendedOperator.MAX_POOL_2D ) ) ) ) and target_vertex['node_type'] in ( ExtendedOperator.CONV_2D, ExtendedOperator.CONV_3D, ExtendedOperator.DEPTHWISE_CONV_2D, ExtendedOperator.MAX_POOL_2D, ) and source_vertex.outdegree() == 1 and target_vertex['op'].padding == Padding.VALID ) def is_slicing_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( target_vertex['node_type'] in (ExtendedOperator.SLICE, ExtendedOperator.STRIDED_SLICE) and ( len(target_vertex['op'].inputs) == 3 or (len(target_vertex['op'].inputs) == 4 and np.all(target_vertex['op'].inputs[3].tensor == 1)) ) and source_vertex['node_type'] in ( ExtendedOperator.TRANSPOSE_CONV, ExtendedOperator.CONV_3D_TRANSPOSE, ) and source_vertex.outdegree() == 1 and source_vertex['op'].padding == Padding.VALID ) def is_requantize_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] in ( ExtendedOperator.FULLY_CONNECTED, ExtendedOperator.GENERIC_CONV, ExtendedOperator.ADD, ExtendedOperator.SUB, ExtendedOperator.MUL, ExtendedOperator.DIV, ExtendedOperator.MAX_POOL_2D, ExtendedOperator.AVERAGE_POOL_2D, ExtendedOperator.GENERIC_DECONV, ) and source_vertex['op'].outputs[0].quantization is not None and target_vertex['node_type'] == ExtendedOperator.QUANTIZE ) def is_activ_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] in ( ExtendedOperator.FULLY_CONNECTED, ExtendedOperator.GENERIC_CONV, ExtendedOperator.ADD, ExtendedOperator.SUB, ExtendedOperator.MUL, ExtendedOperator.DIV, ExtendedOperator.MAX_POOL_2D, ExtendedOperator.AVERAGE_POOL_2D, ExtendedOperator.GENERIC_DECONV, ) and target_vertex['node_type'] in (ExtendedOperator.RELU, ExtendedOperator.RELU6) and source_vertex['op'].fusedActivationFunction == ActivationFunctionType.NONE and source_vertex.outdegree() == 1 ) def is_requantize_node(vertex: ig.Vertex, graph_converter: ig.Graph): return ( vertex['node_type'] == ExtendedOperator.QUANTIZE and vertex['op'].inputs[0].quantization is not None and vertex['op'].outputs[0].quantization is not None ) def is_large_cat_node(vertex: ig.Vertex, graph_converter: ig.Graph): return vertex['node_type'] == ExtendedOperator.CONCATENATION and len(vertex['op'].inputs) > 10 def is_high_dim_transpose_node(vertex: ig.Vertex, graph_converter: ig.Graph, max_transpose_dims: int): return vertex['node_type'] == ExtendedOperator.TRANSPOSE and vertex['op'].inputs[1].tensor.size > max_transpose_dims def is_group_conv_node(vertex: ig.Vertex, graph_converter: ig.Graph): return ( vertex['node_type'] == ExtendedOperator.CONV_2D and vertex['op'].inputs[0].shape[3] != vertex['op'].inputs[1].shape[3] ) def is_group_deconv_node(vertex: ig.Vertex, graph_converter: ig.Graph): return ( vertex['node_type'] == ExtendedOperator.TRANSPOSE_CONV and vertex['op'].outputs[0].shape[3] != vertex['op'].inputs[1].shape[0] ) def is_transformable_node(vertex: ig.Vertex, graph_converter: ig.Graph): return vertex['node_type'] <= ExtendedOperator.BATCH_NORM and vertex.outdegree() >= 1 def is_transformable_transpose_node(vertex: ig.Vertex, graph_converter: ig.Graph): return ( vertex['node_type'] == ExtendedOperator.TRANSPOSE and vertex.outdegree() >= 1 and is_transpose_same_to_reshape_op(vertex['op']) ) def is_multi_output_op_node(vertex: ig.Vertex, graph_converter: ig.Graph): return vertex['node_type'] >= 0 and len(vertex['outputs']) > 1 and vertex.outdegree() > 0 def is_quantize_elementwise_op_edge(edge: ig.Edge, graph_converter: ig.Graph, with_lstm: bool): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( ( source_vertex['node_type'] == ExtendedOperator.DEQUANTIZE and is_quantizable_rewrite_op(target_vertex['node_type'], target_vertex['op'], with_lstm) ) or ( target_vertex['node_type'] == ExtendedOperator.QUANTIZE and is_quantizable_rewrite_op(source_vertex['node_type'], source_vertex['op'], with_lstm) ) ) and target_vertex['op'].inputs[0].name in source_vertex['outputs'] def is_transpose_reshape_op_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( ( source_vertex['node_type'] == ExtendedOperator.TRANSPOSE and target_vertex['node_type'] == ExtendedOperator.RESHAPE ) or ( target_vertex['node_type'] == ExtendedOperator.TRANSPOSE and source_vertex['node_type'] == ExtendedOperator.RESHAPE ) ) and target_vertex['op'].inputs[0].name in source_vertex['outputs'] def is_transpose_elementwise_op_edge(edge: ig.Edge, graph_converter: ig.Graph, quantizable_ops_only: bool): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] if quantizable_ops_only: is_unary = is_elementwise_unary_quantizable_op is_binary = is_elementwise_binary_quantizable_op else: is_unary = is_elementwise_unary_op is_binary = is_elementwise_binary_op return ( ( source_vertex['node_type'] == ExtendedOperator.TRANSPOSE and ( is_unary(target_vertex['node_type'], target_vertex['op']) or is_binary(target_vertex['node_type'], target_vertex['op']) ) ) or ( target_vertex['node_type'] == ExtendedOperator.TRANSPOSE and ( is_unary(source_vertex['node_type'], source_vertex['op']) or is_binary(source_vertex['node_type'], source_vertex['op']) ) ) ) and ( ( target_vertex['node_type'] != ExtendedOperator.SPLIT and target_vertex['op'].inputs[0].name in source_vertex['outputs'] ) or ( target_vertex['node_type'] == ExtendedOperator.SPLIT and target_vertex['op'].inputs[1].name in source_vertex['outputs'] ) ) def is_reshape_elementwise_op_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( ( source_vertex['node_type'] == ExtendedOperator.RESHAPE and ( is_elementwise_unary_op(target_vertex['node_type'], target_vertex['op']) or is_elementwise_binary_op(target_vertex['node_type'], target_vertex['op']) ) ) or ( target_vertex['node_type'] == ExtendedOperator.RESHAPE and ( is_elementwise_unary_op(source_vertex['node_type'], source_vertex['op']) or is_elementwise_binary_op(source_vertex['node_type'], source_vertex['op']) ) ) ) and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name def is_elementwise_reduce_op(op_code: ExtendedOperator, op: tfl.BaseOperator): return ( op_code in ( ExtendedOperator.SUM, ExtendedOperator.ARG_MIN, ExtendedOperator.ARG_MAX, ExtendedOperator.REDUCE_MIN, ExtendedOperator.REDUCE_MAX, ExtendedOperator.REDUCE_PROD, ) and len(op.inputs[0].shape) == len(op.outputs[0].shape) ) or ( op_code == ExtendedOperator.MEAN and len(op.inputs[0].shape) == len(op.outputs[0].shape) and ( len(op.inputs[0].shape) != 4 or ( not np.array_equal(op.inputs[1].tensor, np.array([1, 2], dtype='int32')) and not np.array_equal(op.inputs[1].tensor, np.array([2, 1], dtype='int32')) ) ) ) def is_elementwise_unary_quantizable_op(op_code: ExtendedOperator, op: tfl.BaseOperator): return op_code in ( ExtendedOperator.SOFTMAX, ExtendedOperator.LOG_SOFTMAX, ) def is_elementwise_binary_quantizable_op(op_code: ExtendedOperator, op: tfl.BaseOperator): return False def is_elementwise_unary_op(op_code: ExtendedOperator, op: tfl.BaseOperator): return op_code in ( ExtendedOperator.RELU, ExtendedOperator.SIN, ExtendedOperator.COS, ExtendedOperator.TANH, ExtendedOperator.ELU, ExtendedOperator.PRELU, ExtendedOperator.EXP, ExtendedOperator.LOG, ExtendedOperator.NEG, ExtendedOperator.FLOOR, ExtendedOperator.RELU6, ExtendedOperator.QUANTIZE, ExtendedOperator.DEQUANTIZE, ExtendedOperator.SQRT, ExtendedOperator.RSQRT, ExtendedOperator.CAST, ExtendedOperator.LOGISTIC, ExtendedOperator.HARD_SWISH, ExtendedOperator.LEAKY_RELU, ExtendedOperator.SPLIT, ExtendedOperator.SPLIT_V, ExtendedOperator.UNPACK, ExtendedOperator.PAD, ExtendedOperator.PADV2, ExtendedOperator.MIRROR_PAD, ExtendedOperator.SLICE, ExtendedOperator.STRIDED_SLICE, ExtendedOperator.TILE, ExtendedOperator.GATHER, ExtendedOperator.ABS, ) or is_elementwise_reduce_op(op_code, op) def is_quantizable_rewrite_op(op_code: ExtendedOperator, op: tfl.BaseOperator, with_lstm: bool): return op_code in ( ExtendedOperator.BATCH_MATMUL, ExtendedOperator.SOFTMAX, ExtendedOperator.LOG_SOFTMAX, ExtendedOperator.ABS, ExtendedOperator.SUM, ExtendedOperator.DIV, ExtendedOperator.RSQRT, ExtendedOperator.MAXIMUM, ExtendedOperator.MINIMUM, ) or (with_lstm and op_code == ExtendedOperator.UNIDIRECTIONAL_SEQUENCE_LSTM) def is_elementwise_binary_op(op_code: ExtendedOperator, op: tfl.BaseOperator): return ( op_code in ( ExtendedOperator.CONCATENATION, ExtendedOperator.PACK, ExtendedOperator.ADD, ExtendedOperator.SUB, ExtendedOperator.MUL, ExtendedOperator.DIV, ExtendedOperator.MAXIMUM, ExtendedOperator.MINIMUM, ExtendedOperator.SQUARED_DIFFERENCE, ) and len(op.inputs) >= 2 ) def is_non_passthrough_op(op_code: ExtendedOperator, op: tfl.BaseOperator): return op_code in ( ExtendedOperator.CONV_2D, ExtendedOperator.AVERAGE_POOL_2D, ExtendedOperator.DEPTHWISE_CONV_2D, ExtendedOperator.MAX_POOL_2D, ) def is_ending_with_noop_edge(edge: ig.Edge, graph_converter: ig.Graph, branch: bool = False): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] if branch: source_cond_var = source_vertex.outdegree() >= 1 else: source_cond_var = source_vertex.outdegree() == 1 return ( source_cond_var and target_vertex.outdegree() >= 1 and target_vertex['op'] is not None and target_vertex['op'].inputs[0].name in source_vertex['outputs'] and ( ( target_vertex['node_type'] == ExtendedOperator.RESHAPE and target_vertex['op'].inputs[0].shape == target_vertex['op'].outputs[0].shape ) or ( target_vertex['node_type'] == ExtendedOperator.TRANSPOSE and (np.diff(target_vertex['op'].inputs[1].tensor) == 1).all() ) or ( target_vertex['node_type'] in (ExtendedOperator.PAD, ExtendedOperator.PADV2, ExtendedOperator.MIRROR_PAD) and target_vertex['op'].inputs[0].shape == target_vertex['op'].outputs[0].shape ) or ( target_vertex['node_type'] == ExtendedOperator.TILE and target_vertex['op'].inputs[0].shape == target_vertex['op'].outputs[0].shape ) or ( target_vertex['node_type'] in (ExtendedOperator.SLICE, ExtendedOperator.STRIDED_SLICE) and target_vertex['op'].inputs[0].shape == target_vertex['op'].outputs[0].shape ) or ( target_vertex['node_type'] == ExtendedOperator.CONCATENATION and len(target_vertex['op'].inputs) == 1 and len(target_vertex['op'].outputs) == 1 and target_vertex['op'].inputs[0].shape == target_vertex['op'].outputs[0].shape ) or ( target_vertex['node_type'] == ExtendedOperator.GATHER and target_vertex['op'].inputs[0].shape == target_vertex['op'].outputs[0].shape and (np.diff(target_vertex['op'].inputs[1].tensor) == 1).all() ) or ( target_vertex['node_type'] == ExtendedOperator.CAST and target_vertex['op'].inDataType == target_vertex['op'].outDataType ) ) ) def is_bmm_add_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] out_dim_idx = None if source_vertex['node_type'] == ExtendedOperator.BATCH_MATMUL: out_dim_idx = -1 elif source_vertex['node_type'] == ExtendedOperator.FULLY_CONNECTED: out_dim_idx = 0 return ( out_dim_idx is not None and target_vertex['node_type'] == ExtendedOperator.ADD and source_vertex['op'].inputs[0].tensor.ndim >= 2 and source_vertex['op'].inputs[1].tensor.ndim == 2 and target_vertex['op'].inputs[1].tensor.ndim == 1 and target_vertex['op'].inputs[1].shape[0] == source_vertex['op'].inputs[1].shape[out_dim_idx] and source_vertex.outdegree() == 1 and target_vertex.outdegree() >= 1 and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name ) def is_wrapped_reshape_within_transpose_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( ( ( target_vertex['node_type'] == ExtendedOperator.TRANSPOSE and source_vertex['node_type'] == ExtendedOperator.RESHAPE ) or ( source_vertex['node_type'] == ExtendedOperator.TRANSPOSE and target_vertex['node_type'] == ExtendedOperator.RESHAPE ) ) and source_vertex.outdegree() == 1 and target_vertex.outdegree() >= 1 and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name ) def is_slice_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] in (ExtendedOperator.SLICE, ExtendedOperator.STRIDED_SLICE) and source_vertex.outdegree() == 1 and target_vertex['node_type'] in (ExtendedOperator.SLICE, ExtendedOperator.STRIDED_SLICE) and target_vertex.outdegree() >= 1 and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name and source_vertex['op'].inputs[1].buffer is not None and source_vertex['op'].inputs[2].buffer is not None ) def is_transpose_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.TRANSPOSE and source_vertex.outdegree() == 1 and target_vertex['node_type'] == ExtendedOperator.TRANSPOSE and target_vertex.outdegree() >= 1 and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name ) def is_gather_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.GATHER and source_vertex.outdegree() == 1 and target_vertex['node_type'] == ExtendedOperator.GATHER and target_vertex.outdegree() >= 1 and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name and source_vertex['op'].axis == target_vertex['op'].axis and source_vertex['op'].batchDims == target_vertex['op'].batchDims ) def is_reshape_branch_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.RESHAPE and source_vertex.outdegree() > 1 and target_vertex['node_type'] == ExtendedOperator.RESHAPE and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name ) def is_transpose_branch_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.TRANSPOSE and source_vertex.outdegree() > 1 and target_vertex['node_type'] == ExtendedOperator.TRANSPOSE and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name ) def is_dequant_quant_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph, q_first: bool): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] if q_first: cond = ( source_vertex['node_type'] == ExtendedOperator.QUANTIZE and target_vertex['node_type'] == ExtendedOperator.DEQUANTIZE ) else: cond = ( source_vertex['node_type'] == ExtendedOperator.DEQUANTIZE and target_vertex['node_type'] == ExtendedOperator.QUANTIZE ) return ( cond and source_vertex.outdegree() == 1 and target_vertex.outdegree() >= 1 and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name ) def is_reshape_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.RESHAPE and source_vertex.outdegree() == 1 and target_vertex['node_type'] == ExtendedOperator.RESHAPE and target_vertex.outdegree() >= 1 and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name ) def is_constant_transpose_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.CONSTANT_NODE and target_vertex['node_type'] == ExtendedOperator.TRANSPOSE and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name and target_vertex.outdegree() >= 1 ) def is_constant_reshape_fusable_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.CONSTANT_NODE and target_vertex['node_type'] == ExtendedOperator.RESHAPE and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name and target_vertex.outdegree() >= 1 ) def is_quant_dequant_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.INPUT_NODE and target_vertex['node_type'] == ExtendedOperator.QUANTIZE and source_vertex['outputs'][0] == target_vertex['op'].inputs[0].name ) or ( source_vertex['node_type'] == ExtendedOperator.DEQUANTIZE and target_vertex['node_type'] == ExtendedOperator.OUTPUT_NODE ) def is_transpose_same_to_reshape_op(op: tfl.BaseOperator): num_elements = np.prod(op.inputs[0].shape) input_shape = np.array(op.inputs[0].shape, dtype='int32') output_shape = np.array(op.outputs[0].shape, dtype='int32') if np.array_equal(input_shape[input_shape != 1], output_shape[output_shape != 1]): input_tensor = np.arange(num_elements).reshape(input_shape) perm = op.inputs[1].tensor new_tensor = np.transpose(input_tensor, perm) return np.array_equal(new_tensor.flatten(), input_tensor.flatten()) else: return False def is_conv2d_gather_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.CONV_2D and target_vertex['node_type'] == ExtendedOperator.GATHER and source_vertex.outdegree() == 1 and target_vertex['op'].inputs[1].buffer is not None and target_vertex['op'].axis == 3 and source_vertex['op'].inputs[1].tensor.shape[0] == target_vertex['op'].inputs[1].tensor.shape[0] ) def is_gather_conv2d_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.GATHER and target_vertex['node_type'] == ExtendedOperator.CONV_2D and source_vertex.outdegree() == 1 and source_vertex['op'].inputs[1].buffer is not None and source_vertex['op'].axis == 3 and source_vertex['op'].inputs[1].tensor.shape[0] == target_vertex['op'].inputs[1].tensor.shape[3] ) def is_reciprocal_sqrt_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.SQRT and target_vertex['node_type'] == ExtendedOperator.DIV and source_vertex.outdegree() == 1 ) def is_tile_binary_op_edge(edge: ig.Edge, graph_converter: ig.Graph): source_vertex = graph_converter.vs[edge.source] target_vertex = graph_converter.vs[edge.target] return ( source_vertex['node_type'] == ExtendedOperator.TILE and target_vertex['node_type'] in ( ExtendedOperator.ADD, ExtendedOperator.SUB, ExtendedOperator.MUL, ExtendedOperator.DIV, ) and source_vertex.outdegree() == 1 ) def op_input_dims(op: tfl.BaseOperator): dim_indices = None if isinstance(op, (tfl.ConcatenationOperator, tfl.GatherOperator, tfl.PackOperator, tfl.UnpackOperator)): dim_indices = op.axis elif isinstance(op, tfl.SplitOperator): dim_indices = op.inputs[0].tensor.item() elif isinstance(op, tfl.SplitVOperator): dim_indices = op.inputs[2].tensor.item() elif isinstance(op, (tfl.PadOperator, tfl.Padv2Operator, tfl.MirrorPadOperator)): pads = np.sum(op.inputs[1].tensor, axis=-1) nonzero_idx = np.nonzero(pads)[0] # TODO: support multi indices if nonzero_idx.size == 1: dim_indices = nonzero_idx[0] elif isinstance(op, tfl.PreluOperator): w_shape = np.array(op.inputs[1].shape, dtype='int32') nonzero_idx = np.nonzero(w_shape != 1)[0] if nonzero_idx.size == 1: dim_indices = nonzero_idx[0] + 1 elif isinstance(op, (tfl.SliceOperator, tfl.StridedSliceOperator, tfl.TileOperator)): old_shape = np.array(op.inputs[0].shape) new_shape = np.array(op.outputs[0].shape) diff = new_shape - old_shape nonzero_idx = np.nonzero(diff)[0] # TODO: support multi indices if nonzero_idx.size == 1: dim_indices = nonzero_idx[0] elif isinstance( op, ( tfl.SumOperator, tfl.MeanOperator, tfl.ArgMinOperator, tfl.ArgMaxOperator, tfl.ReduceMinOperator, tfl.ReduceMaxOperator, tfl.ReduceProdOperator, ), ): # TODO: support multi indices if op.inputs[1].tensor.size == 1: dim_indices = op.inputs[1].tensor[0] return dim_indices def op_input_indices(op: tfl.BaseOperator): if isinstance(op, (tfl.ConcatenationOperator, tfl.PackOperator)): input_indices = range(len(op.inputs)) elif isinstance(op, tfl.SplitOperator): input_indices = (1,) elif isinstance(op, (tfl.BatchMatmulOperator, tfl.MinimumOperator, tfl.MaximumOperator)): input_indices = range(2) elif isinstance( op, (tfl.AddOperator, tfl.SubOperator, tfl.MulOperator, tfl.DivOperator, tfl.SquaredDifferenceOperator) ): if len(op.inputs[1].shape) == 1 and op.inputs[1].shape[0] == 1: input_indices = range(1) elif len(op.inputs[0].shape) == 1 and op.inputs[0].shape[0] == 1: input_indices = (1,) else: input_indices = range(2) else: input_indices = range(1) return input_indices def fuse_bn_weight(eps, scale, var, weight, transpose): if transpose: shape = [1, -1] + [1] * (len(weight.shape) - 2) else: shape = [-1, 1] + [1] * (len(weight.shape) - 2) inv = 1 / np.sqrt(var + eps) return weight * (scale * inv).reshape(shape) def fuse_bn_bias(eps, scale, var, mean, bn_b, activ_b): inv = 1 / np.sqrt(var + eps) if activ_b is not None: if activ_b.shape != mean.shape and activ_b.ndim == 1 and activ_b.size == 1: activ_b = activ_b.repeat(mean.size) return (activ_b - mean) * inv * scale + bn_b else: return (-mean) * inv * scale + bn_b def fuse_rev_bn_weight(eps, scale, var, weight): shape = [1, -1] + [1] * (len(weight.shape) - 2) inv = 1 / np.sqrt(var + eps) return weight * (scale * inv).reshape(shape) def fuse_rev_bn_bias(eps, scale, var, mean, bn_b, activ_b, weight): reduced_dims = tuple([i for i in range(len(weight.shape)) if i > 1]) inv = 1 / np.sqrt(var + eps) fused_b = bn_b - mean * inv * scale if weight.shape[1] == 1 and mean.shape[0] > 1: offset_b = (weight.sum(reduced_dims) * fused_b.reshape(-1, 1)).reshape(-1) else: offset_b = np.matmul(weight.sum(reduced_dims), fused_b.reshape(-1, 1)).reshape(-1) if activ_b is not None: if activ_b.shape != mean.shape and activ_b.ndim == 1 and activ_b.size == 1: activ_b = activ_b.repeat(mean.size) return activ_b + offset_b else: return offset_b def fuse_slices(seq: typing.Iterable[ig.Vertex]): cur_start = None cur_end = None cur_strides = None for node in seq: assert node['node_type'] in (ExtendedOperator.SLICE, ExtendedOperator.STRIDED_SLICE) next_start = node['op'].inputs[1].tensor if cur_strides is None: cur_strides = np.ones_like(next_start, dtype='int32') if cur_start is None: cur_start = np.zeros_like(next_start, dtype='int32') if node['node_type'] == ExtendedOperator.SLICE: next_size = node['op'].inputs[2].tensor next_end = cur_start + (next_start + next_size) * cur_strides next_strides = np.ones_like(next_start, dtype='int32') else: next_end = node['op'].inputs[2].tensor next_end = cur_start + next_end * cur_strides next_strides = node['op'].inputs[3].tensor if cur_end is None: cur_start = next_start cur_end = next_end cur_strides = next_strides else: cur_start += next_start * cur_strides cur_end = np.min((cur_end, next_end), axis=0) cur_strides = cur_strides * next_strides return cur_start, cur_end, cur_strides def fuse_transpose_perms(seq: typing.Iterable[ig.Vertex]): cur_perm = None for node in seq: assert node['node_type'] in (ExtendedOperator.TRANSPOSE, ExtendedOperator.GATHER) next_perm = node['op'].inputs[1].tensor if cur_perm is None: cur_perm = next_perm else: cur_perm = cur_perm[next_perm] return cur_perm def fuse_transpose_perms_extended(seq: typing.Iterable[ig.Vertex]): cur_perm = None # Reverse the sequence if dim is expanding if seq[1]['node_type'] == ExtendedOperator.RESHAPE: if len(seq[1]['op'].inputs[0].shape) < len(seq[1]['op'].outputs[0].shape): seq = list(reversed(list(seq))) for node in seq: if node['node_type'] == ExtendedOperator.TRANSPOSE: next_perm = node['op'].inputs[1].tensor if cur_perm is None: cur_perm = next_perm else: cur_perm = cur_perm[next_perm] elif node['node_type'] == ExtendedOperator.RESHAPE: if len(seq[1]['op'].inputs[0].shape) > len(seq[1]['op'].outputs[0].shape): old_shape = node['op'].inputs[0].shape new_shape = node['op'].outputs[0].shape else: new_shape = node['op'].inputs[0].shape old_shape = node['op'].outputs[0].shape if old_shape != new_shape: if len(old_shape) != len(new_shape): new_shape_padded = list(new_shape) + [None] * (len(old_shape) - len(new_shape)) next_perm = [] new_idx = 0 while new_idx < len(new_shape): for old, item in zip(old_shape, cur_perm): if old == new_shape_padded[new_idx] and item not in next_perm: next_perm.append(item) new_idx += 1 cur_perm = np.argsort(next_perm) else: mapping = {} for i in range(len(new_shape)): mapping.setdefault(new_shape[i], []) mapping[new_shape[i]].append(i) next_perm = [0] * len(old_shape) for i in range(len(old_shape)): next_perm[i] = mapping[old_shape[i]].pop(0) cur_perm = cur_perm[next_perm] return cur_perm def fuse_connected_edges( filtered_pairs: typing.List[typing.Iterable[ig.Vertex]], ) -> typing.List[typing.Iterable[ig.Vertex]]: while True: heads = {n[0]: i for i, n in enumerate(filtered_pairs)} tails = {n[-1]: i for i, n in enumerate(filtered_pairs)} connectables = heads.keys() & tails.keys() if len(connectables) > 0: curr_filtered = [] for seq in filtered_pairs: head_connectable = seq[0] in connectables preserve = head_connectable and filtered_pairs[tails[seq[0]]][0] in connectables if preserve: curr_filtered.append(seq) elif not head_connectable: if seq[-1] in connectables: curr_filtered.append(seq + filtered_pairs[heads[seq[-1]]][1:]) else: curr_filtered.append(seq) filtered_pairs = curr_filtered else: break return filtered_pairs def is_simple_reshape(orig_shape, new_shape, mapping: typing.Optional[typing.Dict[int, int]] = None): if orig_shape == new_shape: if mapping is not None: for i in range(len(orig_shape)): mapping[i] = i return True i = 0 j = 0 while True: if i == len(orig_shape) and j == len(new_shape): break elif i == len(orig_shape): if new_shape[j] == 1: j += 1 else: break elif j == len(new_shape): if orig_shape[i] == 1: i += 1 else: break elif orig_shape[i] == new_shape[j]: if mapping is not None: mapping[i] = j i += 1 j += 1 elif orig_shape[i] == 1: i += 1 elif new_shape[j] == 1: j += 1 else: break if i != len(orig_shape) or j != len(new_shape): return False else: return True def reshape_mapping(shape_1, shape_2): i = 0 j = 0 acc_l = 1 start_l = 0 acc_r = 1 start_r = 0 mapping_l = [] mapping_r = [] sign = None while i < len(shape_1) or j < len(shape_2): if i < len(shape_1) and j < len(shape_2): if start_l == i and start_r == j and shape_1[i] == shape_2[j]: mapping_l.append([i]) mapping_r.append([j]) acc_l = 1 acc_r = 1 i += 1 j += 1 start_l = i start_r = j sign = None else: if sign in ('l', None): acc_l = shape_1[i] * acc_l if sign in ('r', None): acc_r = shape_2[j] * acc_r if acc_l == acc_r: mapping_l.append(list(range(start_l, i + 1))) mapping_r.append(list(range(start_r, j + 1))) acc_l = 1 acc_r = 1 i += 1 j += 1 start_l = i start_r = j sign = None elif acc_l < acc_r: sign = 'l' i += 1 else: sign = 'r' j += 1 elif i < len(shape_1): assert shape_1[i] == 1 mapping_l[-1].append(i) i += 1 else: assert shape_2[j] == 1 mapping_r[-1].append(j) j += 1 non_one_mapping_l = [] non_one_mapping_r = [] for ml, mr in zip(mapping_l, mapping_r): new_ml = [i for i in ml if shape_1[i] != 1] new_mr = [j for j in mr if shape_2[j] != 1] if len(new_ml) > 0 and len(new_mr) > 0: non_one_mapping_l.append(new_ml) non_one_mapping_r.append(new_mr) return mapping_l, mapping_r, non_one_mapping_l, non_one_mapping_r def elinimate_sequences( graph_converter: CommonGraph, filtered_pairs: typing.List[typing.Iterable[ig.Vertex]], remove_first_pred: typing.Union[bool, typing.Callable] = False, remove_first_node_action: typing.Optional[typing.Callable] = None, remove_last_pred: typing.Union[bool, typing.Callable] = True, remove_last_node_action: typing.Optional[typing.Callable] = None, skip_pred: typing.Union[bool, typing.Callable] = False, input_idx: int = 0, force_forward_input: bool = False, ): remove_ids = [] actions = [] for seq in filtered_pairs: first_node = seq[0] last_node = seq[-1] if type(skip_pred) is bool: skip = skip_pred elif skip_pred is not None: skip = skip_pred(seq) if skip: continue if type(remove_first_pred) is bool: remove_first = remove_first_pred custom_data = None elif remove_first_pred is not None: remove_first, custom_data = remove_first_pred(seq) if type(remove_last_pred) is bool: remove_last = remove_last_pred custom_data_last = None elif remove_last_pred is not None: remove_last, custom_data_last = remove_last_pred(seq) # If the first node can also be eliminated, then set the previous node as the first node if remove_first: first_node = graph_converter.graph.vs.find( name=graph_converter.tensor_node_map[first_node['op'].inputs[input_idx].name] ) if not remove_last: last_node = seq[-2] output_idx = 0 if first_node == seq[0]: next_idx = 1 else: next_idx = 0 output_name = seq[next_idx]['op'].inputs[input_idx].name output_idx = first_node['outputs'].index(output_name) # We use the forward input tensor under the following circumstances. # 1. If the previous node before the sequence is an input node # 2. If the first node has multiple outputs and the last node doesn't connect to output nodes use_forward_input = False if first_node['node_type'] == ExtendedOperator.INPUT_NODE: use_forward_input = True branch = first_node.outdegree() > 1 has_output_nodes = False for edge in last_node.out_edges(): target_vertex = edge.target_vertex if target_vertex['node_type'] in (ExtendedOperator.OUTPUT_NODE, ExtendedOperator.UNUSED_NODE): if use_forward_input: # Cannot optimize away ops between i/o nodes skip = True else: has_output_nodes = True break if branch: output_outdegree = 0 for edge in first_node.out_edges(): target_vertex = edge.target_vertex if target_vertex == seq[next_idx]: continue if target_vertex['node_type'] in (ExtendedOperator.OUTPUT_NODE, ExtendedOperator.UNUSED_NODE): if has_output_nodes and edge['label'] == output_name: output_outdegree += 1 break else: names = [t.name for t in target_vertex['op'].inputs] if output_name in names: output_outdegree += 1 break if not has_output_nodes: use_forward_input = True elif output_outdegree > 0: skip = True if force_forward_input and not use_forward_input: if not has_output_nodes: use_forward_input = True else: skip = True if skip: continue if use_forward_input: # Find out the output of the first node in the sequence new_output = first_node['outputs'][output_idx] assert new_output in graph_converter.tensor_map # For each node that is next of the last node, we connect it with the first node # Also, the replace the tensors when needed graph_converter.replace_next_tensors(last_node, first_node, new_output) else: # Find out the output of the last node in the sequence new_output = last_node['outputs'][0] assert new_output in graph_converter.tensor_map # For each node that is next of the last node, we connect it with the first node graph_converter.connect_next_tensors(last_node, first_node, new_output) # Update graph, prepare to drop the output tensor of the intermediate nodes and use the output tensor of # the last node instead first_node['outputs'][output_idx] = new_output if first_node['op'] is not None: first_node['op'].outputs[output_idx] = graph_converter.tensor_map[new_output] graph_converter.tensor_node_map[new_output] = first_node['name'] # When the first node is a constant node, we need to set the buffer back if first_node['node_type'] == ExtendedOperator.CONSTANT_NODE and not use_forward_input: if seq[0]['node_type'] == ExtendedOperator.CONSTANT_NODE: old_tensor = graph_converter.tensor_map[seq[0]['name']] else: old_tensor = seq[0]['op'].inputs[input_idx] new_tensor = seq[-1]['op'].outputs[0] new_tensor.buffer = old_tensor.buffer if remove_first and remove_last: # Push the sequence to the removing list remove_ids.extend([x.index for x in seq]) else: # Collect actions when removing the first node start_index = 0 end_index = len(seq) if not remove_first: start_index = 1 if remove_first_node_action is not None: action = remove_first_node_action(first_node, last_node, custom_data) if action is not None: actions.extend(action) if not remove_last: end_index = len(seq) - 1 if remove_last_node_action is not None: action = remove_last_node_action(first_node, last_node, custom_data_last) if action is not None: actions.extend(action) # Push the sequence (except the first node) to the removing list remove_ids.extend([x.index for x in seq[start_index:end_index]]) for func, args in actions: func(*args) graph_converter.graph.delete_vertices(remove_ids) def expand_op_outputs_in_branches( nodes: typing.List[ig.Vertex], new_op_func: typing.Callable[[ig.Vertex, ig.Vertex, ig.Vertex], None], graph_converter: CommonGraph, ): actions = [] for node in nodes: preserve_node = None prev_node_name = node['op'].inputs[0].name prev_node = graph_converter.graph.vs.find(name=graph_converter.tensor_node_map[prev_node_name]) # Collect next nodes and choose one to preserve next_nodes = [] for edge in node.out_edges(): next_node = graph_converter.graph.vs[edge.target] if preserve_node is None or next_node['node_type'] == ExtendedOperator.OUTPUT_NODE: preserve_node = next_node next_nodes.append(next_node) # For the filtered nodes, use the cloned op as the previous op filtered_nodes = list(set(next_nodes) - set([preserve_node])) for next_node in filtered_nodes: actions.extend(new_op_func(node, prev_node, next_node)) # Process actions for func, args in actions: node = args[0] func(*args) def get_same_padding_args(input_shape, filter_shape, strides, dilation): dim = len(input_shape) padding = [0] * dim for i in range(dim): if input_shape[i] % strides[i] == 0: padding[i] = max(1 - strides[i] + (filter_shape[i] - 1) * dilation[i], 0) else: padding[i] = max(1 + (filter_shape[i] - 1) * dilation[i] - (input_shape[i] % strides[i]), 0) pad_args = [[0, 0]] + [[x // 2, x - x // 2] for x in padding] + [[0, 0]] return pad_args