From 49f0610447fa3f111f263768fba121a12cce5d54 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 4 May 2020 14:06:57 -0700 Subject: [PATCH] Add options --disable_layer_norm, --disable_gelu and --enable_gelu_approximation (#3750) --- .../python/tools/bert/BertOnnxModel.py | 29 +++++++++++++++++-- onnxruntime/python/tools/bert/README.md | 2 +- .../tools/bert/bert_model_optimization.py | 24 +++++++++++++++ 3 files changed, 52 insertions(+), 3 deletions(-) diff --git a/onnxruntime/python/tools/bert/BertOnnxModel.py b/onnxruntime/python/tools/bert/BertOnnxModel.py index 97df569082..28ac282f14 100644 --- a/onnxruntime/python/tools/bert/BertOnnxModel.py +++ b/onnxruntime/python/tools/bert/BertOnnxModel.py @@ -18,11 +18,14 @@ logger = logging.getLogger(__name__) class BertOptimizationOptions: def __init__(self, model_type): + self.enable_gelu = True + self.enable_layer_norm = True self.enable_attention = True self.enable_skip_layer_norm = True self.enable_embed_layer_norm = True self.enable_bias_skip_layer_norm = True self.enable_bias_gelu = True + self.enable_gelu_approximation = False if model_type == 'gpt2': self.enable_skip_layer_norm = False @@ -515,6 +518,23 @@ class BertOnnxModel(OnnxModel): self.remove_nodes(nodes_to_remove) self.add_nodes(nodes_to_add) + def gelu_approximation(self): + nodes_to_add = [] + nodes_to_remove = self.get_nodes_by_op_type("Gelu") + self.get_nodes_by_op_type("BiasGelu") + for node in nodes_to_remove: + new_node = onnx.helper.make_node("FastGelu", + inputs=node.input, + outputs=node.output, + name=self.create_node_name("FastGelu", node.op_type + "_Approximation")) + new_node.domain = "com.microsoft" + nodes_to_add.append(new_node) + + if len(nodes_to_add) > 0: + logger.info(f"Gelu approximation count:{len(nodes_to_add)}") + + self.remove_nodes(nodes_to_remove) + self.add_nodes(nodes_to_add) + def fuse_add_bias_skip_layer_norm(self): logger.debug(f"start Bias and SkipLayerNormalization fusion...") input_name_to_nodes = self.input_name_to_nodes() @@ -1129,9 +1149,11 @@ class BertOnnxModel(OnnxModel): self.prune_graph() def optimize(self, options: BertOptimizationOptions = None): - self.fuse_layer_norm() + if (options is None) or options.enable_layer_norm: + self.fuse_layer_norm() - self.fuse_gelu() + if (options is None) or options.enable_gelu: + self.fuse_gelu() self.preprocess() @@ -1159,6 +1181,9 @@ class BertOnnxModel(OnnxModel): # Fuse SkipLayerNormalization and Add Bias before it. self.fuse_add_bias_skip_layer_norm() + if (options is not None and options.enable_gelu_approximation): + self.gelu_approximation() + self.remove_unused_constant() # Use symbolic batch dimension in input and output. diff --git a/onnxruntime/python/tools/bert/README.md b/onnxruntime/python/tools/bert/README.md index e1fc9a15a0..5399fa08fc 100644 --- a/onnxruntime/python/tools/bert/README.md +++ b/onnxruntime/python/tools/bert/README.md @@ -55,7 +55,7 @@ See below for description of some options of bert_model_optimization.py: - **input**: input model path - **output**: output model path - **model_type**: (*defaul: bert*) - There are 4 model types: *bert* (exported by PyTorch), *bert_tf* (BERT exported by PyTorch), *bert_keras* (BERT exported by keras2onnx) and *gpt2* (exported by PyTorch) respectively. + There are 4 model types: *bert* (exported by PyTorch), *bert_tf* (BERT exported by tf2onnx), *bert_keras* (BERT exported by keras2onnx) and *gpt2* (exported by PyTorch) respectively. - **num_heads**: (*default: 12*) Number of attention heads. BERT-base and BERT-large has 12 and 16 respectively. - **hidden_size**: (*default: 768*) diff --git a/onnxruntime/python/tools/bert/bert_model_optimization.py b/onnxruntime/python/tools/bert/bert_model_optimization.py index 17673c2247..64c74916fc 100644 --- a/onnxruntime/python/tools/bert/bert_model_optimization.py +++ b/onnxruntime/python/tools/bert/bert_model_optimization.py @@ -156,6 +156,24 @@ def parse_arguments(): help="disable Add Bias and Gelu/FastGelu fusion") parser.set_defaults(disable_bias_gelu=False) + parser.add_argument('--disable_layer_norm', + required=False, + action='store_true', + help="disable LayerNormalization fusion") + parser.set_defaults(disable_layer_norm=False) + + parser.add_argument('--disable_gelu', + required=False, + action='store_true', + help="disable Gelu fusion") + parser.set_defaults(disable_gelu=False) + + parser.add_argument('--enable_gelu_approximation', + required=False, + action='store_true', + help="enable Gelu/BiasGelu to FastGelu conversion") + parser.set_defaults(enable_gelu_approximation=False) + parser.add_argument('--verbose', required=False, action='store_true') parser.set_defaults(verbose=False) @@ -173,6 +191,10 @@ def parse_arguments(): def get_optimization_options(args): optimization_options = BertOptimizationOptions(args.model_type) + if args.disable_gelu: + optimization_options.enable_gelu = False + if args.disable_layer_norm: + optimization_options.enable_layer_norm = False if args.disable_attention: optimization_options.enable_attention = False if args.disable_skip_layer_norm: @@ -183,6 +205,8 @@ def get_optimization_options(args): optimization_options.enable_bias_skip_layer_norm = False if args.disable_bias_gelu: optimization_options.enable_bias_gelu = False + if args.enable_gelu_approximation: + optimization_options.enable_gelu_approximation = True return optimization_options