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' ] self.elementwise_layers = [ 'fluid.layers.elementwise_add', 'fluid.layers.elementwise_sub', 'fluid.layers.elementwise_mul', 'fluid.layers.elementwise_div' ] 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 strip_direct_layers(self, graph): # 构建opt_graph # 删除所有direct_layers， 便于对transpose进行优化 opt_graph = copy.deepcopy(graph) remove_layer_ids = set() for layer_id, layer in opt_graph.layers.items(): if layer.kernel in self.direct_layers: layer_out = opt_graph.edges_out[layer_id] layer_in = opt_graph.edges_in[layer_id] if len(layer_out) == 0 or len(layer_in) == 0: continue assert len( layer_in ) == 1, "There should be only 1 input for direct layers." remove_layer_ids.add(layer_id) for layer_id in remove_layer_ids: opt_graph.del_layer(layer_id) return opt_graph def run(self, graph): optimized_transpose_layers = list() modified_layer_attrs = dict() modified_parameters = dict() scanned_layers = set() total_layer_num = len(graph.layers) 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_layer_ids = list() elementwise_layer_ids = list() concat_layer_ids = list() can_be_optimized = True modified_attrs = dict() parameter_layers = list() parameters = dict() for out in _graph.edges_out[layer_id]: if _graph.layers[out].kernel == "fluid.layers.transpose": if _graph.layers[out].attrs["perm"] != [0, 3, 1, 2]: can_be_optimized = False continue transpose_layer_ids.append(out) elif _graph.layers[out].kernel in self.elementwise_layers: elementwise_layer_ids.append(out) elif _graph.layers[out].kernel == "fluid.layers.concat": elementwise_layer_ids.append(out) concat_layer_ids.append(out) else: can_be_optimized = False break visited_layers = set() while len(elementwise_layer_ids) > 0 and can_be_optimized: current_id = elementwise_layer_ids.pop(0) visited_layers.add(current_id) for out in _graph.edges_out[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 if out not in visited_layers: transpose_layer_ids.append(out) elif _graph.layers[ out].kernel in self.elementwise_layers: if out not in visited_layers: elementwise_layer_ids.append(out) elif _graph.layers[out].kernel == "fluid.layers.concat": if out not in visited_layers: elementwise_layer_ids.append(out) concat_layer_ids.append(out) else: can_be_optimized = False break all_create_parameter = True for ipt in _graph.edges_in.get(current_id, []): if _graph.layers[ ipt].kernel == "fluid.layers.transpose": all_creater_parameter = False if _graph.layers[ipt].attrs["perm"] != [0, 2, 3, 1]: can_be_optimized = False break if ipt not in visited_layers: transpose_layer_ids.append(ipt) elif _graph.layers[ ipt].kernel in self.elementwise_layers: all_creater_parameter = False if ipt not in visited_layers: elementwise_layer_ids.append(ipt) elif _graph.layers[ipt].kernel == "fluid.layers.concat": all_creater_parameter = False if ipt not in visited_layers: elementwise_layer_ids.append(ipt) concat_layer_ids.append(ipt) elif _graph.layers[ ipt].kernel == "fluid.layers.create_parameter": if ipt not in visited_layers: elementwise_layer_ids.append(ipt) parameter_layers.append(ipt) else: can_be_optimized = False break if all_create_parameter: can_be_optimized = False break if not can_be_optimized: break if not can_be_optimized: continue concat_layer_ids = list(set(concat_layer_ids)) for l in concat_layer_ids: axis = _graph.layers[l].attrs.get('axis', 0) _graph.layers[l].attrs['axis'] = [0, 2, 3, 1][axis] modified_attrs[l] = _graph.layers[l].attrs parameter_layers = list(set(parameter_layers)) for l in parameter_layers: for o in _graph.edges_out[l]: if _graph.layers[o].kernel in self.elementwise_layers: axis = _graph.layers[o].attrs.get('axis', -1) _graph.layers[o].attrs['axis'] = [0, 3, 1, 2][axis] modified_attrs[o] = _graph.layers[o].attrs else: can_be_optimized = False break if not can_be_optimized: break s = _graph.layers[l].attrs['shape'] p = _graph.parameters[_graph.layers[l].outputs[0]] if len(s) == 4: _graph.layers[l].attrs[ 'shape'] = [s[0], s[3], s[1], s[2]] modified_attrs[l] = _graph.layers[l].attrs parameters[_graph.layers[l].outputs[0]] = np.transpose( p, (0, 3, 1, 2)) elif len(s) == 3: _graph.layers[l].attrs['shape'] = [s[2], s[0], s[1]] modified_attrs[l] = _graph.layers[l].attrs parameters[_graph.layers[l].outputs[0]] = np.transpose( p, (2, 0, 1)) if not can_be_optimized: continue transpose_layer_ids.append(layer_id) transpose_layer_ids = list(set(transpose_layer_ids)) for transpose_layer_id in transpose_layer_ids: _graph.del_layer(transpose_layer_id) optimized_transpose_layers.extend(transpose_layer_ids) modified_layer_attrs.update(modified_attrs) modified_parameters.update(parameters) return True return False before_transpose_num = self.get_transpose_num(graph) opt_graph = self.strip_direct_layers(graph) total_layer_num = len(opt_graph.layers) while strip_transpose(opt_graph): pass for layer_id in optimized_transpose_layers: graph.del_layer(layer_id) for layer_id, attrs in modified_layer_attrs.items(): graph.layers[layer_id].attrs = attrs for name, parameter in modified_parameters.items(): graph.parameters[name] = parameter 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))