import copy import sys class TransposeOpt: def __init__(self): self.image_layers = [ 'fluid.layers.conv2d', 'fluid.layers.batch_norm', 'fluid.layers.conv2d_transpose', 'fluid.layers.resize_nearest', 'fluid.layers.resize_bilinear', 'fluid.layers.pool2d', 'fluid.layers.pad2d' ] self.direct_layers = [ 'fluid.layers.relu', 'fluid.layers.relu6', 'fluid.layers.abs', 'fluid.layers.sigmoid', 'fluid.layers.exp', 'fluid.layers.rsqrt', 'fluid.layers.swish_f32', 'fluid.layers.tanh', 'fluid.layers.softplus', 'fluid.layers.leaky_relu', 'fluid.layers.floor', 'fluid.layers.erf', 'fluid.layers.swish' ] self.elementwise_layers = [ 'fluid.layers.elementwise_add', 'fluid.layers.elementwise_sub', 'fluid.layers.elementwise_mul', 'fluid.layers.elementwise_div' ] self.reduce_layers = [ 'fluid.layers.reduce_mean', 'fluid.layers.reduce_all', 'fluid.layers.reduce_max', 'fluid.layers.reduce_any', 'fluid.layers.reduce_sum', 'fluid.layers.reduce_prod' ] def get_transpose_num(self, graph): count = 0 for layer_id, layer in graph.layers.items(): if layer.kernel == "fluid.layers.transpose": count += 1 return count def run(self, graph): print("Optimize: TransposeOpt...") total_layer_num = len(graph.layers) scanned_layers = set() optimized_transpose_layers = list() optimized_reduce_layers = list() optimized_concat_layers = list() optimized_elementwise_layers = list() def strip_transpose(_graph): layers = copy.deepcopy(_graph.layers) for layer_id, layer in layers.items(): if layer_id in scanned_layers: continue scanned_layers.add(layer_id) percent = round(len(scanned_layers) / total_layer_num * 100, 2) sys.stderr.write("\rOptimize Transpose Layers...{}%".format( percent)) if layer.kernel != "fluid.layers.transpose": continue if layer.attrs["perm"] != [0, 2, 3, 1]: continue transpose_layers = list() propagate_layers = list() reduce_layers = list() concat_layers = list() # 此elementwise_layers专用于存储shape(4) + shape(1)的形式layer elementwise_layers = list() can_be_optimized = True for out in _graph.edges_out.get(layer_id, []): if _graph.layers[out].kernel == "fluid.layers.transpose": if _graph.layers[out].attrs["perm"] != [0, 3, 1, 2]: can_be_optimized = False break transpose_layers.append(out) elif _graph.layers[out].kernel in self.elementwise_layers: propagate_layers.append(out) elif _graph.layers[out].kernel in self.direct_layers: if _graph.layers[out].outputs[0] in _graph.outputs: can_be_optimized = False break propagate_layers.append(out) elif _graph.layers[out].kernel in self.reduce_layers: if _graph.layers[out].outputs[0] in _graph.outputs: can_be_optimized = False break if not _graph.layers[out].attrs.get('keep_dim', False): can_be_optimized = False break propagate_layers.append(out) reduce_layers.append(out) elif _graph.layers[out].kernel == "fluid.layers.concat": if _graph.layers[out].outputs[0] in _graph.outputs: can_be_optimized = False break propagate_layers.append(out) concat_layers.append(out) else: can_be_optimized = False break visited_layers = set() while len(propagate_layers) > 0 and can_be_optimized: current_id = propagate_layers.pop(0) visited_layers.add(current_id) for out in _graph.edges_out.get(current_id, []): if _graph.layers[ out].kernel == "fluid.layers.transpose": if _graph.layers[out].attrs["perm"] != [0, 3, 1, 2]: can_be_optimized = False break transpose_layers.append(out) elif _graph.layers[ out].kernel in self.elementwise_layers: if _graph.layers[out].outputs[0] in _graph.outputs: can_be_optimized = False break if out not in visited_layers: propagate_layers.append(out) elif _graph.layers[out].kernel in self.direct_layers: if _graph.layers[out].outputs[0] in _graph.outputs: can_be_optimized = False break if out not in visited_layers: propagate_layers.append(out) elif _graph.layers[out].kernel in self.reduce_layers: if _graph.layers[out].outputs[0] in _graph.outputs: can_be_optimized = False break if not _graph.layers[out].attrs.get('keep_dim', False): can_be_optimized = False break if out not in visited_layers: propagate_layers.append(out) reduce_layers.append(out) elif _graph.layers[out].kernel == "fluid.layers.concat": if _graph.layers[out].outputs[0] in _graph.outputs: can_be_optimized = False break if out not in visited_layers: propagate_layers.append(out) concat_layers.append(out) else: can_be_optimized = False break for ipt in _graph.edges_in.get(current_id, []): if _graph.layers[ current_id].kernel in self.elementwise_layers: try: x_shape = _graph.layers[ current_id].input_shapes['x'] y_shape = _graph.layers[ current_id].input_shapes['y'] if _graph.layers[ipt].outputs[ 0] == _graph.layers[current_id].inputs[ 'x']: if len(x_shape) <= 1: elementwise_layers.append(current_id) continue elif _graph.layers[ipt].outputs[ 0] == _graph.layers[current_id].inputs[ 'y']: if len(y_shape) <= 1: elementwise_layers.append(current_id) continue else: raise Exception( "Unexcepted situation happend while optimizing transpose" ) except Exception as e: can_be_optimized = False break if _graph.layers[ ipt].kernel == "fluid.layers.transpose": if _graph.layers[ipt].attrs["perm"] != [0, 2, 3, 1]: can_be_optimized = False break if ipt not in visited_layers: transpose_layers.append(ipt) elif _graph.layers[ ipt].kernel in self.elementwise_layers: if _graph.layers[ipt].outputs[0] in _graph.outputs: can_be_optimized = False break if ipt not in visited_layers: propagate_layers.append(ipt) elif _graph.layers[ipt].kernel in self.direct_layers: if _graph.layers[ipt].outputs[0] in _graph.outputs: can_be_optimized = False break if ipt not in visited_layers: propagate_layers.append(ipt) elif _graph.layers[ipt].kernel in self.reduce_layers: if _graph.layers[ipt].outputs[0] in _graph.outputs: can_be_optimized = False break if not _graph.layers[ipt].attrs.get('keep_dim', False): can_be_optimized = False break if ipt not in visited_layers: propagate_layers.append(ipt) reduce_layers.append(ipt) elif _graph.layers[ipt].kernel == "fluid.layers.concat": if _graph.layers[ipt].outputs[0] in _graph.outputs: can_be_optimized = False break if ipt not in visited_layers: propagate_layers.append(ipt) concat_layers.append(ipt) else: can_be_optimized = False break if not can_be_optimized: break if not can_be_optimized: continue transpose_layers.append(layer_id) transpose_layers = list(set(transpose_layers)) for l in transpose_layers: if graph.layers[l].outputs[0] in graph.outputs: can_be_optimized = False break if not can_be_optimized: continue for l in transpose_layers: _graph.del_layer(l) optimized_transpose_layers.extend(transpose_layers) optimized_reduce_layers.extend(reduce_layers) optimized_concat_layers.extend(concat_layers) optimized_elementwise_layers.extend(elementwise_layers) return True return False before_transpose_num = self.get_transpose_num(graph) opt_graph = copy.deepcopy(graph) total_layer_num = len(opt_graph.layers) while strip_transpose(opt_graph): pass for layer_id in list(set(optimized_transpose_layers)): graph.del_layer(layer_id) for layer_id in list(set(optimized_reduce_layers)): dim = graph.layers[layer_id].attrs.get('dim', None) if dim is not None: for i in range(len(dim)): dim[i] = [0, 2, 3, 1][dim[i]] graph.layers[layer_id].attrs['dim'] = dim for layer_id in list(set(optimized_concat_layers)): axis = graph.layers[layer_id].attrs.get('axis', 0) graph.layers[layer_id].attrs['axis'] = [0, 2, 3, 1][axis] for layer_id in list(set(optimized_elementwise_layers)): axis = graph.layers[layer_id].attrs.get('axis', -1) graph.layers[layer_id].attrs['axis'] = [0, 2, 3, 1][axis] current_transpose_num = self.get_transpose_num(graph) print( "\nTranspose layers optimized, before: transpose_num={}, after: transpose_num={}". format(before_transpose_num, current_transpose_num))