mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Transformer benchmark: add option to use raw attention mask (#4446)
* Update benchmark and optimizer to add an option to use raw attention mask * Remove temporary model in optimizer
This commit is contained in:
Tianlei Wu
GitHub
parent
b156ae4448
commit
05757b4c3c
6 changed files with 106 additions and 44 deletions
|
|
@ -235,11 +235,13 @@ def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwri
|
|||
|
||||
|
||||
def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_attention_heads, hidden_size, use_gpu,
|
||||
fp16, overwrite):
|
||||
fp16, use_raw_attention_mask, overwrite):
|
||||
if overwrite or not os.path.exists(optimized_model_path):
|
||||
from optimizer import optimize_model
|
||||
from onnx_model_bert import BertOptimizationOptions
|
||||
optimization_options = BertOptimizationOptions(model_type)
|
||||
if use_raw_attention_mask:
|
||||
optimization_options.use_raw_attention_mask()
|
||||
if fp16:
|
||||
optimization_options.enable_gelu_approximation = True
|
||||
|
||||
|
|
@ -264,7 +266,7 @@ def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_a
|
|||
|
||||
|
||||
def export_onnx_model(model_name, cache_dir, onnx_dir, input_names, use_gpu, fp16, optimize_onnx, validate_onnx,
|
||||
overwrite):
|
||||
use_raw_attention_mask, overwrite):
|
||||
config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir)
|
||||
model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir)
|
||||
model.cpu()
|
||||
|
|
@ -310,7 +312,7 @@ def export_onnx_model(model_name, cache_dir, onnx_dir, input_names, use_gpu, fp1
|
|||
model_type = MODELS[model_name][3]
|
||||
optimized_model_path = get_onnx_file_path(onnx_dir, model_name, len(input_names), True, use_gpu, fp16, False)
|
||||
optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, config.num_attention_heads,
|
||||
config.hidden_size, use_gpu, fp16, overwrite)
|
||||
config.hidden_size, use_gpu, fp16, use_raw_attention_mask, overwrite)
|
||||
|
||||
onnx_model_path = optimized_model_path
|
||||
if validate_onnx:
|
||||
|
|
@ -390,7 +392,8 @@ def allocateOutputBuffers(output_buffers, max_last_state_size, max_pooler_size,
|
|||
|
||||
|
||||
def run_onnxruntime(use_gpu, model_names, fp16, batch_sizes, sequence_lengths, repeat_times, input_counts,
|
||||
optimize_onnx, validate_onnx, cache_dir, onnx_dir, verbose, overwrite, disable_ort_io_binding):
|
||||
optimize_onnx, validate_onnx, cache_dir, onnx_dir, verbose, overwrite, disable_ort_io_binding,
|
||||
use_raw_attention_mask):
|
||||
import onnxruntime
|
||||
|
||||
results = []
|
||||
|
|
@ -414,7 +417,7 @@ def run_onnxruntime(use_gpu, model_names, fp16, batch_sizes, sequence_lengths, r
|
|||
with torch.no_grad():
|
||||
onnx_model_file, is_valid_onnx_model, vocab_size, max_sequence_length = export_onnx_model(
|
||||
model_name, cache_dir, onnx_dir, input_names, use_gpu, fp16, optimize_onnx, validate_onnx,
|
||||
overwrite)
|
||||
use_raw_attention_mask, overwrite)
|
||||
if not is_valid_onnx_model:
|
||||
continue
|
||||
|
||||
|
|
@ -691,6 +694,12 @@ def parse_arguments():
|
|||
help='Disable running ONNX Runtime with binded inputs and outputs. ')
|
||||
parser.set_defaults(disable_ort_io_binding=False)
|
||||
|
||||
parser.add_argument('--use_raw_attention_mask',
|
||||
required=False,
|
||||
action='store_true',
|
||||
help='Use raw attention mask in Attention operator for Bert models.')
|
||||
parser.set_defaults(use_raw_attention_mask=False)
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
|
@ -742,7 +751,7 @@ def main():
|
|||
results += run_onnxruntime(args.use_gpu, args.models, args.fp16, args.batch_sizes, args.sequence_lengths,
|
||||
args.test_times, args.input_counts, args.optimize_onnx, args.validate_onnx,
|
||||
args.cache_dir, args.onnx_dir, args.verbose, args.overwrite,
|
||||
args.disable_ort_io_binding)
|
||||
args.disable_ort_io_binding, args.use_raw_attention_mask)
|
||||
except:
|
||||
logger.error(f"Exception", exc_info=True)
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
#--------------------------------------------------------------------------
|
||||
import numpy as np
|
||||
from logging import getLogger
|
||||
from enum import Enum
|
||||
from onnx import helper, numpy_helper, TensorProto
|
||||
from onnx_model import OnnxModel
|
||||
from fusion_base import Fusion
|
||||
|
|
@ -12,6 +13,12 @@ from fusion_utils import FusionUtils
|
|||
logger = getLogger(__name__)
|
||||
|
||||
|
||||
class AttentionMaskFormat:
|
||||
MaskIndexEnd = 0
|
||||
MaskIndexEndAndStart = 1
|
||||
AttentionMask = 2
|
||||
|
||||
|
||||
class AttentionMask():
|
||||
"""
|
||||
Fuse Attention subgraph into one Attention node.
|
||||
|
|
@ -23,6 +30,10 @@ class AttentionMask():
|
|||
# A lookup table with mask input as key, and cast (to int32) output as value
|
||||
self.mask_casted = {}
|
||||
self.utils = FusionUtils(model)
|
||||
self.mask_format = AttentionMaskFormat.MaskIndexEnd
|
||||
|
||||
def set_mask_format(self, mask_format: AttentionMaskFormat):
|
||||
self.mask_format = mask_format
|
||||
|
||||
def set_mask_indice(self, mask, mask_index):
|
||||
if mask in self.mask_indice:
|
||||
|
|
@ -47,7 +58,12 @@ class AttentionMask():
|
|||
if casted:
|
||||
self.mask_casted[input] = input_name
|
||||
|
||||
# Add a mask processing node
|
||||
# Attention supports int32 attention mask (2D) since 1.4.0
|
||||
if self.mask_format == AttentionMaskFormat.AttentionMask:
|
||||
self.mask_indice[input] = input_name
|
||||
return input_name
|
||||
|
||||
# Add a mask processing node to convert attention mask to mask index (1D)
|
||||
output_name = self.model.create_node_name('mask_index')
|
||||
mask_index_node = helper.make_node('ReduceSum',
|
||||
inputs=[input_name],
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ class FusionEmbedLayerNoMask(Fusion):
|
|||
self.utils = FusionUtils(model)
|
||||
|
||||
def fuse(self, node, input_name_to_nodes, output_name_to_node):
|
||||
# already fused. Assumes that only one mebedding layer in a transformer model.
|
||||
# Embed layer has been fused. Assumes that only one mebedding layer in a transformer model.
|
||||
if self.nodes_to_add:
|
||||
return
|
||||
|
||||
|
|
@ -61,7 +61,7 @@ class FusionEmbedLayerNoMask(Fusion):
|
|||
normalize_node = node
|
||||
word_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 0])
|
||||
if word_embedding_path is None:
|
||||
logger.info("Failed to find word embedding")
|
||||
logger.info("Word embedding path is not found. Embed layer cannot be fused.")
|
||||
return
|
||||
add_node, word_embedding_gather = word_embedding_path
|
||||
input_ids = word_embedding_gather.input[1]
|
||||
|
|
@ -87,7 +87,7 @@ class FusionEmbedLayerNoMask(Fusion):
|
|||
if position_embedding_path is not None:
|
||||
position_embedding_weight_node, position_embedding_expand = position_embedding_path
|
||||
else:
|
||||
logger.info("Failed to find position embedding")
|
||||
logger.info("Position embedding path is not found. Embed layer cannot be fused.")
|
||||
return
|
||||
|
||||
if position_embedding_shape is not None and position_embedding_shape.input[0] != input_ids:
|
||||
|
|
@ -98,7 +98,7 @@ class FusionEmbedLayerNoMask(Fusion):
|
|||
if segment_embedding_path is None:
|
||||
segment_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 1])
|
||||
if segment_embedding_path is None:
|
||||
logger.info("Failed to find segment embedding")
|
||||
logger.info("Segment embedding is not found. Embed layer cannot be fused.")
|
||||
return
|
||||
_, segment_embedding_gather = segment_embedding_path
|
||||
else:
|
||||
|
|
@ -206,19 +206,23 @@ class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask):
|
|||
else:
|
||||
embed_node = self.nodes_to_add.pop()
|
||||
mask_input_name = next(iter(self.mask_indice))
|
||||
mask_output_name = self.mask_indice[mask_input_name]
|
||||
mask_node = output_name_to_node[mask_output_name]
|
||||
|
||||
self.nodes_to_remove.extend([mask_node])
|
||||
|
||||
# store inputs for further processing
|
||||
self.mask_input_name = mask_input_name
|
||||
|
||||
# When mask has been casted to int32, use that casted one as input of embed layer norm.
|
||||
if mask_input_name in self.mask_casted:
|
||||
mask_input_name = self.mask_casted[mask_input_name]
|
||||
mask_output_name = self.mask_indice[mask_input_name]
|
||||
mask_node = output_name_to_node[mask_output_name]
|
||||
# Remove mask processing node (like ReduceSum), but keep Cast node for 2D attention mask.
|
||||
if mask_node.op_type == "ReduceSum" or mask_node.op_type == "MaskIndex":
|
||||
self.nodes_to_remove.extend([mask_node])
|
||||
|
||||
# When mask has been casted to int32, use that casted one as input of embed layer norm.
|
||||
if mask_input_name in self.mask_casted:
|
||||
mask_input_name = self.mask_casted[mask_input_name]
|
||||
|
||||
embed_node.input.append(mask_input_name)
|
||||
|
||||
embed_node.output[1] = mask_output_name
|
||||
|
||||
embed_node.input.append(mask_input_name)
|
||||
embed_node.output[1] = mask_output_name
|
||||
self.nodes_to_add.append(embed_node)
|
||||
self.prune_graph = True
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from fusion_reshape import FusionReshape
|
|||
from fusion_layernorm import FusionLayerNormalization, FusionLayerNormalizationTF
|
||||
from fusion_skiplayernorm import FusionSkipLayerNormalization, FusionBiasSkipLayerNormalization
|
||||
from fusion_embedlayer import FusionEmbedLayerNormalization
|
||||
from fusion_attention import FusionAttention, AttentionMask
|
||||
from fusion_attention import FusionAttention, AttentionMask, AttentionMaskFormat
|
||||
from fusion_gelu import FusionGelu
|
||||
from fusion_fastgelu import FusionFastGelu
|
||||
from fusion_biasgelu import FusionBiasGelu
|
||||
|
|
@ -30,9 +30,14 @@ class BertOptimizationOptions:
|
|||
self.enable_bias_skip_layer_norm = True
|
||||
self.enable_bias_gelu = True
|
||||
self.enable_gelu_approximation = False
|
||||
self.attention_mask_format = AttentionMaskFormat.MaskIndexEnd
|
||||
|
||||
if model_type == 'gpt2':
|
||||
self.enable_skip_layer_norm = False
|
||||
self.attention_mask_format = AttentionMaskFormat.AttentionMask
|
||||
|
||||
def use_raw_attention_mask(self):
|
||||
self.attention_mask_format = AttentionMaskFormat.AttentionMask
|
||||
|
||||
|
||||
class BertOnnxModel(OnnxModel):
|
||||
|
|
@ -158,8 +163,9 @@ class BertOnnxModel(OnnxModel):
|
|||
# After:
|
||||
# input_ids --> Shape --> ConstantOfShape -->Cast --> EmbedLayerNormaliation/ReduceSum
|
||||
# TODO: merge ConstantOfShape -->Cast to ConstantOfShape (need update the data type of value)
|
||||
if node.op_type == 'EmbedLayerNormalization' or node.op_type == 'ReduceSum':
|
||||
i = 1 if node.op_type == 'EmbedLayerNormalization' else 0
|
||||
op_input_id = {"EmbedLayerNormalization": 1, "ReduceSum": 0, "Attention": 3}
|
||||
if node.op_type in op_input_id:
|
||||
i = op_input_id[node.op_type]
|
||||
parent_nodes = self.match_parent_path(
|
||||
node, ['Cast', 'ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [i, 0, 0, 0, 0, 0],
|
||||
output_name_to_node)
|
||||
|
|
@ -208,6 +214,8 @@ class BertOnnxModel(OnnxModel):
|
|||
self.fuse_skip_layer_norm()
|
||||
|
||||
if (options is None) or options.enable_attention:
|
||||
if options is not None:
|
||||
self.attention_mask.set_mask_format(options.attention_mask_format)
|
||||
self.fuse_attention()
|
||||
|
||||
if (options is None) or options.enable_embed_layer_norm:
|
||||
|
|
|
|||
|
|
@ -87,7 +87,7 @@ def optimize_by_onnxruntime(onnx_model_path: str,
|
|||
assert 'CUDAExecutionProvider' in session.get_providers() # Make sure there is GPU
|
||||
|
||||
assert os.path.exists(optimized_model_path) and os.path.isfile(optimized_model_path)
|
||||
logger.info("Save optimized model by onnxruntime to {}".format(optimized_model_path))
|
||||
logger.debug("Save optimized model by onnxruntime to {}".format(optimized_model_path))
|
||||
return optimized_model_path
|
||||
|
||||
|
||||
|
|
@ -186,6 +186,12 @@ def _parse_arguments():
|
|||
help="enable Gelu/BiasGelu to FastGelu conversion")
|
||||
parser.set_defaults(enable_gelu_approximation=False)
|
||||
|
||||
parser.add_argument('--use_raw_attention_mask',
|
||||
required=False,
|
||||
action='store_true',
|
||||
help="use raw attention mask instead of mask index in attention operator")
|
||||
parser.set_defaults(use_raw_attention_mask=False)
|
||||
|
||||
parser.add_argument('--verbose', required=False, action='store_true')
|
||||
parser.set_defaults(verbose=False)
|
||||
|
||||
|
|
@ -225,6 +231,9 @@ def _get_optimization_options(args):
|
|||
optimization_options.enable_bias_gelu = False
|
||||
if args.enable_gelu_approximation:
|
||||
optimization_options.enable_gelu_approximation = True
|
||||
if args.use_raw_attention_mask:
|
||||
optimization_options.use_raw_attention_mask()
|
||||
|
||||
return optimization_options
|
||||
|
||||
|
||||
|
|
@ -258,16 +267,14 @@ def optimize_model(input,
|
|||
"""
|
||||
(optimizer_class, producer, run_onnxruntime) = MODEL_CLASSES[model_type]
|
||||
|
||||
input_model_path = input
|
||||
|
||||
if opt_level > 1: # Optimization specified for an execution provider.
|
||||
input_model_path = optimize_by_onnxruntime(input_model_path, use_gpu=use_gpu, opt_level=opt_level)
|
||||
temp_model_path = optimize_by_onnxruntime(input, use_gpu=use_gpu, opt_level=opt_level)
|
||||
elif run_onnxruntime:
|
||||
# Use Onnxruntime to do optimizations (like constant folding and cast elimation) that is not specified to exection provider.
|
||||
# CPU provider is used here so that there is no extra node for GPU memory copy.
|
||||
input_model_path = optimize_by_onnxruntime(input_model_path, use_gpu=False, opt_level=1)
|
||||
temp_model_path = optimize_by_onnxruntime(input, use_gpu=False, opt_level=1)
|
||||
|
||||
model = load_model(input_model_path, format=None, load_external_data=True)
|
||||
model = load_model(temp_model_path, format=None, load_external_data=True)
|
||||
|
||||
if model.producer_name and producer != model.producer_name:
|
||||
logger.warning(
|
||||
|
|
@ -282,6 +289,10 @@ def optimize_model(input,
|
|||
if not only_onnxruntime:
|
||||
optimizer.optimize(optimization_options)
|
||||
|
||||
# Remove the temporary model.
|
||||
os.remove(temp_model_path)
|
||||
logger.debug("Remove tempoary model: {}".format(temp_model_path))
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ if [ "$run_cpu" = true ] ; then
|
|||
average_over=100
|
||||
fi
|
||||
|
||||
# enable optimizer (use script instead of OnnxRuntime for graph optimization)
|
||||
# Enable optimizer (use script instead of OnnxRuntime for graph optimization)
|
||||
use_optimizer=true
|
||||
|
||||
# Batch Sizes and Sequence Lengths
|
||||
|
|
@ -45,12 +45,21 @@ models_to_test="bert-base-cased roberta-base gpt2"
|
|||
# If you have mutliple GPUs, you can choose one GPU for test. Here is an example to use the second GPU:
|
||||
# export CUDA_VISIBLE_DEVICES=1
|
||||
|
||||
#This script will generate a logs file with a list of commands used in tests.
|
||||
# This script will generate a logs file with a list of commands used in tests.
|
||||
echo echo "ort=$run_ort torch=$run_torch torchscript=$run_torchscript gpu_fp32=$run_gpu_fp32 gpu_fp16=$run_gpu_fp16 cpu=$run_cpu optimizer=$use_optimizer batch=$batch_sizes sequence=$sequence_length models=$models_to_test" >> benchmark.log
|
||||
|
||||
#Set it to false to skip testing. You can use it to dry run this script with the log file.
|
||||
# Set it to false to skip testing. You can use it to dry run this script with the log file.
|
||||
run_tests=true
|
||||
|
||||
# Directory for downloading pretrained models.
|
||||
cache_dir="./cache_models"
|
||||
|
||||
# Directory for ONNX models
|
||||
onnx_dir="./onnx_models"
|
||||
|
||||
# Use raw attention mask in Attention operator or not.
|
||||
use_raw_attention_mask=false
|
||||
|
||||
# -------------------------------------------
|
||||
if [ "$run_cpu" = true ] ; then
|
||||
if [ "$run_gpu_fp32" = true ] ; then
|
||||
|
|
@ -79,41 +88,46 @@ fi
|
|||
|
||||
if [ "$run_cli" = true ] ; then
|
||||
echo "Use onnxruntime_tools.transformers.benchmark"
|
||||
optimizer_script="-m onnxruntime_tools.transformers.benchmark"
|
||||
benchmark_script="-m onnxruntime_tools.transformers.benchmark"
|
||||
else
|
||||
optimizer_script="benchmark.py"
|
||||
benchmark_script="benchmark.py"
|
||||
fi
|
||||
|
||||
onnx_export_options="-i $input_counts -v -b 0 --overwrite -f fusion.csv"
|
||||
benchmark_options="-b $batch_sizes -s $sequence_lengths -t $average_over -f fusion.csv -r result.csv -d detail.csv"
|
||||
onnx_export_options="-i $input_counts -v -b 0 --overwrite -f fusion.csv -c $cache_dir --onnx_dir $onnx_dir"
|
||||
benchmark_options="-b $batch_sizes -s $sequence_lengths -t $average_over -f fusion.csv -r result.csv -d detail.csv -c $cache_dir --onnx_dir $onnx_dir"
|
||||
|
||||
if [ "$use_optimizer" = true ] ; then
|
||||
onnx_export_options="$onnx_export_options -o"
|
||||
benchmark_options="$benchmark_options -o"
|
||||
fi
|
||||
|
||||
if [ "$use_raw_attention_mask" = true ] ; then
|
||||
onnx_export_options="$onnx_export_options --use_raw_attention_mask"
|
||||
benchmark_options="$benchmark_options --use_raw_attention_mask"
|
||||
fi
|
||||
|
||||
# -------------------------------------------
|
||||
run_one_test() {
|
||||
if [ "$run_ort" = true ] ; then
|
||||
echo python $optimizer_script -m $1 $onnx_export_options $2 $3 >> benchmark.log
|
||||
echo python $optimizer_script -m $1 $benchmark_options $2 $3 -i $input_counts >> benchmark.log
|
||||
echo python $benchmark_script -m $1 $onnx_export_options $2 $3 >> benchmark.log
|
||||
echo python $benchmark_script -m $1 $benchmark_options $2 $3 -i $input_counts >> benchmark.log
|
||||
if [ "$run_tests" = true ] ; then
|
||||
python $optimizer_script -m $1 $onnx_export_options $2 $3
|
||||
python $optimizer_script -m $1 $benchmark_options $2 $3 -i $input_counts
|
||||
python $benchmark_script -m $1 $onnx_export_options $2 $3
|
||||
python $benchmark_script -m $1 $benchmark_options $2 $3 -i $input_counts
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "$run_torch" = true ] ; then
|
||||
echo python $optimizer_script -e torch -m $1 $benchmark_options $2 $3 >> benchmark.log
|
||||
echo python $benchmark_script -e torch -m $1 $benchmark_options $2 $3 >> benchmark.log
|
||||
if [ "$run_tests" = true ] ; then
|
||||
python $optimizer_script -e torch -m $1 $benchmark_options $2 $3
|
||||
python $benchmark_script -e torch -m $1 $benchmark_options $2 $3
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "$run_torchscript" = true ] ; then
|
||||
echo python $optimizer_script -e torchscript -m $1 $benchmark_options $2 $3 >> benchmark.log
|
||||
echo python $benchmark_script -e torchscript -m $1 $benchmark_options $2 $3 >> benchmark.log
|
||||
if [ "$run_tests" = true ] ; then
|
||||
python $optimizer_script -e torchscript -m $1 $benchmark_options $2 $3
|
||||
python $benchmark_script -e torchscript -m $1 $benchmark_options $2 $3
|
||||
fi
|
||||
fi
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue