codekingpro/portable-devtools
114k
1# Copyright (c) Microsoft Corporation. All rights reserved.
2# Copyright 2018 The HuggingFace Inc. team.
3# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16"""Benchmarking the inference of pretrained transformer models.
17PyTorch/TorchScript benchmark is based on https://github.com/huggingface/transformers/blob/master/examples/benchmarks.py.
18One difference is that random input_ids is generated in this benchmark.
19
20For onnxruntime, this script will convert a pretrained model to ONNX, and optimize it when -o parameter is used.
21
22Example commands:
23 Export all models to ONNX, optimize and validate them:
24 python benchmark.py -b 0 -o -v -i 1 2 3
25 Run OnnxRuntime on GPU for all models:
26 python benchmark.py -g
27 Run OnnxRuntime on GPU for all models with fp32 optimization:
28 python benchmark.py -g -o
29 Run OnnxRuntime on GPU with fp16 optimization:
30 python benchmark.py -g -o -p "fp16"
31 Run TorchScript on GPU for all models:
32 python benchmark.py -e torchscript -g
33 Run TorchScript on GPU for all models with fp16:
34 python benchmark.py -e torchscript -g -p "fp16"
35 Run ONNXRuntime and TorchScript on CPU for all models with quantization:
36 python benchmark.py -e torchscript onnxruntime -p "int8" -o
37 Run OnnxRuntime with bfloat16 fastmath mode kernels on aarch64 platforms with bfloat16 support:
38 python benchmark.py --enable_arm64_bfloat16_fastmath_mlas_gemm
39
40It is recommended to use run_benchmark.sh to launch benchmark.
41"""
42
43import argparse
44import logging
45import os
46import random
47import timeit
48from datetime import datetime
49
50import numpy
51import psutil
52from benchmark_helper import (
53 ConfigModifier,
54 OptimizerInfo,
55 Precision,
56 create_onnxruntime_session,
57 get_latency_result,
58 inference_ort,
59 inference_ort_with_io_binding,
60 output_details,
61 output_fusion_statistics,
62 output_summary,
63 setup_logger,
64)
65from fusion_options import FusionOptions
66from huggingface_models import MODEL_CLASSES, MODELS
67from onnx_exporter import (
68 create_onnxruntime_input,
69 export_onnx_model_from_pt,
70 export_onnx_model_from_tf,
71 load_pretrained_model,
72)
73from packaging import version
74from quantize_helper import QuantizeHelper
75
76logger = logging.getLogger("")
77
78cpu_count = psutil.cpu_count(logical=False)
79
80# Set OMP environment variable before importing onnxruntime or torch.
81if "OMP_NUM_THREADS" not in os.environ:
82 os.environ["OMP_NUM_THREADS"] = str(cpu_count)
83
84import torch # noqa: E402
85from transformers import AutoConfig, AutoTokenizer, LxmertConfig # noqa: E402
86
87
88def run_onnxruntime(
89 use_gpu,
90 provider,
91 model_names,
92 model_class,
93 config_modifier,
94 precision,
95 num_threads,
96 batch_sizes,
97 sequence_lengths,
98 repeat_times,
99 input_counts,
100 optimizer_info,
101 validate_onnx,
102 cache_dir,
103 onnx_dir,
104 verbose,
105 overwrite,
106 disable_ort_io_binding,
107 use_raw_attention_mask,
108 model_fusion_statistics,
109 model_source,
110 enable_arm64_bfloat16_fastmath_mlas_gemm,
111 args,
112):
113 import onnxruntime # noqa: PLC0415
114
115 results = []
116 if (
117 use_gpu
118 and ("CUDAExecutionProvider" not in onnxruntime.get_available_providers())
119 and ("MIGraphXExecutionProvider" not in onnxruntime.get_available_providers())
120 and ("DmlExecutionProvider" not in onnxruntime.get_available_providers())
121 ):
122 logger.error(
123 "Please install onnxruntime-gpu or onnxruntime-directml package instead of onnxruntime, and use a machine with GPU for testing gpu performance."
124 )
125 return results
126
127 warm_up_repeat = 0
128 if provider == "tensorrt":
129 optimizer_info = OptimizerInfo.NOOPT
130 warm_up_repeat = 5
131 if "TensorrtExecutionProvider" not in onnxruntime.get_available_providers():
132 logger.error(
133 "Please install onnxruntime-gpu-tensorrt package, and use a machine with GPU for testing gpu performance."
134 )
135 return results
136
137 if optimizer_info == OptimizerInfo.NOOPT:
138 logger.warning(
139 f"OptimizerInfo is set to {optimizer_info}, graph optimizations specified in FusionOptions are not applied."
140 )
141
142 for model_name in model_names:
143 all_input_names = MODELS[model_name][0]
144 for num_inputs in input_counts:
145 if num_inputs > len(all_input_names):
146 break
147
148 input_names = all_input_names[:num_inputs]
149 args.model_type = MODELS[model_name][3]
150 fusion_options = FusionOptions.parse(args)
151
152 if "pt" in model_source:
153 with torch.no_grad():
154 (
155 onnx_model_file,
156 is_valid_onnx_model,
157 vocab_size,
158 max_sequence_length,
159 ) = export_onnx_model_from_pt(
160 model_name,
161 MODELS[model_name][1],
162 MODELS[model_name][2],
163 MODELS[model_name][3],
164 model_class,
165 config_modifier,
166 cache_dir,
167 onnx_dir,
168 input_names,
169 use_gpu,
170 precision,
171 optimizer_info,
172 validate_onnx,
173 use_raw_attention_mask,
174 overwrite,
175 model_fusion_statistics,
176 fusion_options,
177 )
178 if "tf" in model_source:
179 (
180 onnx_model_file,
181 is_valid_onnx_model,
182 vocab_size,
183 max_sequence_length,
184 ) = export_onnx_model_from_tf(
185 model_name,
186 MODELS[model_name][1],
187 MODELS[model_name][2],
188 MODELS[model_name][3],
189 model_class,
190 config_modifier,
191 cache_dir,
192 onnx_dir,
193 input_names,
194 use_gpu,
195 precision,
196 optimizer_info,
197 validate_onnx,
198 use_raw_attention_mask,
199 overwrite,
200 model_fusion_statistics,
201 fusion_options,
202 )
203
204 if not is_valid_onnx_model:
205 continue
206
207 ort_session = create_onnxruntime_session(
208 onnx_model_file,
209 use_gpu,
210 provider,
211 enable_all_optimization=True,
212 num_threads=num_threads,
213 verbose=verbose,
214 enable_mlas_gemm_fastmath_arm64_bfloat16=enable_arm64_bfloat16_fastmath_mlas_gemm,
215 )
216 if ort_session is None:
217 continue
218
219 ort_output_names = [node_arg.name for node_arg in ort_session.get_outputs()]
220 output_buffers = []
221 device = "cuda" if use_gpu else "cpu"
222 config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir)
223 max_last_state_size = numpy.prod(
224 [
225 max(batch_sizes),
226 max(sequence_lengths),
227 max(vocab_size, config.hidden_size),
228 ]
229 )
230 max_pooler_size = numpy.prod([max(batch_sizes), config.hidden_size])
231 for batch_size in batch_sizes:
232 if batch_size <= 0:
233 continue
234 for sequence_length in sequence_lengths:
235 if max_sequence_length is not None and sequence_length > max_sequence_length:
236 continue
237
238 input_value_type = numpy.int64 if "pt" in model_source else numpy.int32
239 ort_inputs = create_onnxruntime_input(
240 vocab_size,
241 batch_size,
242 sequence_length,
243 input_names,
244 config,
245 input_value_type,
246 )
247 result_template = {
248 "engine": "onnxruntime",
249 "version": onnxruntime.__version__,
250 "providers": provider,
251 "device": device,
252 "optimizer": optimizer_info,
253 "precision": precision,
254 "io_binding": not disable_ort_io_binding,
255 "model_name": model_name,
256 "inputs": num_inputs,
257 "threads": num_threads,
258 "batch_size": batch_size,
259 "sequence_length": sequence_length,
260 "custom_layer_num": config_modifier.get_layer_num(),
261 "datetime": str(datetime.now()),
262 }
263
264 if config.model_type in ["vit", "swin"]:
265 logger.info(
266 f"Run onnxruntime on {model_name} with input shape {[batch_size, 3, config.image_size, config.image_size]}"
267 )
268 else:
269 logger.info(f"Run onnxruntime on {model_name} with input shape {[batch_size, sequence_length]}")
270
271 if disable_ort_io_binding:
272 result = inference_ort(
273 ort_session,
274 ort_inputs,
275 result_template,
276 repeat_times,
277 batch_size,
278 warm_up_repeat,
279 )
280 else:
281 # Get output sizes from a dummy ort run
282 ort_outputs = ort_session.run(ort_output_names, ort_inputs)
283 output_buffer_max_sizes = [max_last_state_size]
284 for i in range(len(ort_outputs)):
285 if i == 2 and MODELS[model_name][3] == "gpt":
286 # past state output max size
287 output_buffer_max_sizes.append(max_pooler_size)
288 else:
289 output_buffer_max_sizes.append(max_last_state_size)
290
291 data_type = numpy.longlong if "pt" in model_source else numpy.intc
292 result = inference_ort_with_io_binding(
293 ort_session,
294 ort_inputs,
295 result_template,
296 repeat_times,
297 ort_output_names,
298 ort_outputs,
299 output_buffers,
300 output_buffer_max_sizes,
301 batch_size,
302 device,
303 data_type,
304 warm_up_repeat,
305 )
306 logger.info(result)
307 results.append(result)
308
309 return results
310
311
312def run_pytorch(
313 use_gpu,
314 model_names,
315 model_class,
316 config_modifier,
317 precision,
318 num_threads,
319 batch_sizes,
320 sequence_lengths,
321 repeat_times,
322 torchscript,
323 torch2,
324 cache_dir,
325 verbose,
326):
327 results = []
328 if use_gpu and not torch.cuda.is_available():
329 logger.error("Please install PyTorch with Cuda, and use a machine with GPU for testing gpu performance.")
330 return results
331
332 torch.set_grad_enabled(False)
333
334 for model_name in model_names:
335 config = AutoConfig.from_pretrained(model_name, torchscript=torchscript, cache_dir=cache_dir)
336 config_modifier.modify(config)
337 model = load_pretrained_model(
338 model_name,
339 config=config,
340 cache_dir=cache_dir,
341 custom_model_class=model_class,
342 )
343
344 if config.model_type in ["vit", "swin"]:
345 # These models don't use sequence lengths, so just pick the first sequence length so that the summary still works
346 sequence_lengths = [sequence_lengths[0]]
347 else:
348 tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
349
350 max_input_size = tokenizer.model_max_length
351
352 logger.debug(f"Model {model}")
353 logger.debug(f"Number of parameters {model.num_parameters()}")
354
355 if precision == Precision.FLOAT16:
356 model.half()
357
358 device = torch.device("cuda:0" if use_gpu else "cpu")
359 model.to(device)
360
361 if precision == Precision.INT8:
362 model = QuantizeHelper.quantize_torch_model(model)
363
364 for batch_size in batch_sizes:
365 if batch_size <= 0:
366 continue
367
368 for sequence_length in sequence_lengths:
369 if config.model_type in ["vit", "swin"]:
370 logger.info(
371 f"Run PyTorch on {model_name} with input shape {[batch_size, 3, config.image_size, config.image_size]}"
372 )
373 input_ids = torch.randn(
374 size=(batch_size, 3, config.image_size, config.image_size),
375 dtype=torch.float16 if precision == Precision.FLOAT16 else torch.float32,
376 device=device,
377 )
378 else:
379 if max_input_size is not None and sequence_length > max_input_size:
380 continue
381
382 logger.info(f"Run PyTorch on {model_name} with input shape {[batch_size, sequence_length]}")
383 input_ids = torch.randint(
384 low=0,
385 high=config.vocab_size - 1,
386 size=(batch_size, sequence_length),
387 dtype=torch.long,
388 device=device,
389 )
390 try:
391 inference = (
392 torch.jit.trace(model, input_ids) if torchscript else torch.compile(model) if torch2 else model
393 )
394 inference(input_ids)
395
396 runtimes = timeit.repeat(lambda: inference(input_ids), repeat=repeat_times, number=1) # noqa: B023
397
398 result = {
399 "engine": "torchscript" if torchscript else "torch2" if torch2 else "torch",
400 "version": torch.__version__,
401 "providers": "NA",
402 "device": "cuda" if use_gpu else "cpu",
403 "optimizer": "",
404 "precision": precision,
405 "io_binding": "",
406 "model_name": model_name,
407 "inputs": 1,
408 "threads": num_threads,
409 "batch_size": batch_size,
410 "sequence_length": sequence_length,
411 "custom_layer_num": config_modifier.get_layer_num(),
412 "datetime": str(datetime.now()),
413 }
414 result.update(get_latency_result(runtimes, batch_size))
415 logger.info(result)
416 results.append(result)
417 except RuntimeError as e:
418 logger.exception(e)
419 torch.cuda.empty_cache()
420
421 return results
422
423
424def run_with_tf_optimizations(do_eager_mode: bool, use_xla: bool):
425 from functools import wraps # noqa: PLC0415
426
427 import tensorflow as tf # noqa: PLC0415
428
429 def run_func(func):
430 @wraps(func)
431 def run_in_eager_mode(*args, **kwargs):
432 return func(*args, **kwargs)
433
434 @wraps(func)
435 @tf.function(jit_compile=use_xla)
436 def run_in_graph_mode(*args, **kwargs):
437 return func(*args, **kwargs)
438
439 if do_eager_mode is True:
440 assert use_xla is False, (
441 "Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`."
442 )
443 return run_in_eager_mode
444 else:
445 return run_in_graph_mode
446
447 return run_func
448
449
450def run_tensorflow(
451 use_gpu,
452 model_names,
453 model_class,
454 config_modifier,
455 precision,
456 num_threads,
457 batch_sizes,
458 sequence_lengths,
459 repeat_times,
460 cache_dir,
461 verbose,
462):
463 results = []
464
465 import tensorflow as tf # noqa: PLC0415
466
467 tf.config.threading.set_intra_op_parallelism_threads(num_threads)
468
469 if not use_gpu:
470 tf.config.set_visible_devices([], "GPU")
471
472 if use_gpu and not tf.test.is_built_with_cuda():
473 logger.error("Please install Tensorflow-gpu, and use a machine with GPU for testing gpu performance.")
474 return results
475
476 if use_gpu: # Restrict TensorFlow to only use the first GPU
477 physical_devices = tf.config.list_physical_devices("GPU")
478 try:
479 tf.config.set_visible_devices(physical_devices[0], "GPU")
480 tf.config.experimental.set_memory_growth(physical_devices[0], True)
481 tf.distribute.OneDeviceStrategy(device="/gpu:0")
482 except RuntimeError as e:
483 logger.exception(e)
484
485 if precision == Precision.FLOAT16 or precision == Precision.INT8:
486 raise NotImplementedError("Mixed precision is currently not supported.")
487
488 for model_name in model_names:
489 config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir)
490 config_modifier.modify(config)
491
492 model = load_pretrained_model(
493 model_name,
494 config=config,
495 cache_dir=cache_dir,
496 custom_model_class=model_class,
497 is_tf_model=True,
498 )
499
500 tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
501
502 max_input_size = tokenizer.model_max_length
503
504 # Define tf.function-decorated forward functions once per model, outside the
505 # batch_size/sequence_length loops. Passing input_ids as an argument (instead
506 # of closing over it) allows tf.function to cache traced graphs by input shape
507 # rather than retracing on every loop iteration. See issue #14953.
508 @run_with_tf_optimizations(do_eager_mode=False, use_xla=False)
509 def encoder_forward(input_ids):
510 return model(input_ids, training=False) # noqa: B023
511
512 @run_with_tf_optimizations(do_eager_mode=False, use_xla=False)
513 def encoder_decoder_forward(input_ids):
514 return model(input_ids, decoder_input_ids=input_ids, training=False) # noqa: B023
515
516 @run_with_tf_optimizations(do_eager_mode=False, use_xla=False)
517 def lxmert_forward(input_ids):
518 feats = tf.random.normal([1, 1, config.visual_feat_dim]) # noqa: B023
519 pos = tf.random.normal([1, 1, config.visual_pos_dim]) # noqa: B023
520 return model( # noqa: B023
521 input_ids,
522 visual_feats=feats,
523 visual_pos=pos,
524 training=False,
525 )
526
527 if config.is_encoder_decoder:
528 inference = encoder_decoder_forward
529 elif isinstance(config, LxmertConfig):
530 inference = lxmert_forward
531 else:
532 inference = encoder_forward
533
534 for batch_size in batch_sizes:
535 if batch_size <= 0:
536 continue
537
538 for sequence_length in sequence_lengths:
539 if max_input_size is not None and sequence_length > max_input_size:
540 continue
541
542 logger.info(f"Run Tensorflow on {model_name} with input shape {[batch_size, sequence_length]}")
543
544 rng = random.Random()
545 values = [rng.randint(0, config.vocab_size - 1) for i in range(batch_size * sequence_length)]
546 input_ids = tf.constant(values, shape=(batch_size, sequence_length), dtype=tf.int32)
547
548 try:
549 inference(input_ids)
550
551 runtimes = timeit.repeat(lambda: inference(input_ids), repeat=repeat_times, number=1) # noqa: B023
552
553 result = {
554 "engine": "tensorflow",
555 "version": tf.__version__,
556 "providers": "NA",
557 "device": "cuda" if use_gpu else "cpu",
558 "optimizer": "",
559 "precision": precision,
560 "io_binding": "",
561 "model_name": model_name,
562 "inputs": 1,
563 "threads": num_threads,
564 "batch_size": batch_size,
565 "sequence_length": sequence_length,
566 "custom_layer_num": config_modifier.get_layer_num(),
567 "datetime": str(datetime.now()),
568 }
569 result.update(get_latency_result(runtimes, batch_size))
570 logger.info(result)
571 results.append(result)
572 except RuntimeError as e:
573 logger.exception(e)
574 from numba import cuda # noqa: PLC0415
575
576 device = cuda.get_current_device()
577 device.reset()
578
579 return results
580
581
582def parse_arguments():
583 parser = argparse.ArgumentParser()
584
585 parser.add_argument(
586 "-m",
587 "--models",
588 required=False,
589 nargs="+",
590 type=str,
591 default=["bert-base-cased", "roberta-base", "gpt2"],
592 choices=list(MODELS.keys()),
593 help="Pre-trained models in the list: " + ", ".join(MODELS.keys()),
594 )
595
596 parser.add_argument(
597 "--model_source",
598 required=False,
599 nargs=1,
600 type=str,
601 default="pt",
602 choices=["pt", "tf"],
603 help="Export onnx from pt or tf",
604 )
605
606 parser.add_argument(
607 "--model_class",
608 required=False,
609 type=str,
610 default=None,
611 choices=list(MODEL_CLASSES),
612 help="Model type selected in the list: " + ", ".join(MODEL_CLASSES),
613 )
614
615 parser.add_argument(
616 "-e",
617 "--engines",
618 required=False,
619 nargs="+",
620 type=str,
621 default=["onnxruntime"],
622 choices=["onnxruntime", "torch", "torch2", "torchscript", "tensorflow"],
623 help="Engines to benchmark",
624 )
625
626 parser.add_argument(
627 "-c",
628 "--cache_dir",
629 required=False,
630 type=str,
631 default=os.path.join(".", "cache_models"),
632 help="Directory to cache pre-trained models",
633 )
634
635 parser.add_argument(
636 "--onnx_dir",
637 required=False,
638 type=str,
639 default=os.path.join(".", "onnx_models"),
640 help="Directory to store onnx models",
641 )
642
643 parser.add_argument("-g", "--use_gpu", required=False, action="store_true", help="Run on gpu device")
644
645 parser.add_argument(
646 "--provider",
647 required=False,
648 type=str,
649 default=None,
650 help="Execution provider to use",
651 )
652
653 parser.add_argument(
654 "-p",
655 "--precision",
656 type=Precision,
657 default=Precision.FLOAT32,
658 choices=list(Precision),
659 help="Precision of model to run. fp32 for full precision, fp16 for half precision, and int8 for quantization",
660 )
661
662 parser.add_argument("--verbose", required=False, action="store_true", help="Print more information")
663
664 parser.add_argument(
665 "--overwrite",
666 required=False,
667 action="store_true",
668 help="Overwrite existing models",
669 )
670
671 parser.add_argument(
672 "-o",
673 "--optimizer_info",
674 type=OptimizerInfo,
675 default=OptimizerInfo.BYSCRIPT,
676 choices=list(OptimizerInfo),
677 help="Optimizer info: Use optimizer.py to optimize onnx model as default. Can also choose from by_ort and no_opt",
678 )
679
680 parser.add_argument(
681 "-v",
682 "--validate_onnx",
683 required=False,
684 action="store_true",
685 help="Validate ONNX model",
686 )
687
688 parser.add_argument(
689 "-f",
690 "--fusion_csv",
691 required=False,
692 default=None,
693 help="CSV file for saving summary results of graph optimization.",
694 )
695
696 parser.add_argument(
697 "-d",
698 "--detail_csv",
699 required=False,
700 default=None,
701 help="CSV file for saving detail results.",
702 )
703
704 parser.add_argument(
705 "-r",
706 "--result_csv",
707 required=False,
708 default=None,
709 help="CSV file for saving summary results.",
710 )
711
712 parser.add_argument(
713 "-i",
714 "--input_counts",
715 required=False,
716 nargs="+",
717 default=[1],
718 type=int,
719 choices=[1, 2, 3],
720 help="Number of ONNX model inputs. Please use 1 for fair comparison with Torch or TorchScript.",
721 )
722
723 parser.add_argument(
724 "-t",
725 "--test_times",
726 required=False,
727 default=100,
728 type=int,
729 help="Number of repeat times to get average inference latency.",
730 )
731
732 parser.add_argument("-b", "--batch_sizes", nargs="+", type=int, default=[1])
733
734 parser.add_argument(
735 "-s",
736 "--sequence_lengths",
737 nargs="+",
738 type=int,
739 default=[4, 8, 16, 32, 64, 128, 256],
740 )
741
742 parser.add_argument(
743 "--disable_ort_io_binding",
744 required=False,
745 action="store_true",
746 help="Disable running ONNX Runtime with binded inputs and outputs. ",
747 )
748 parser.set_defaults(disable_ort_io_binding=False)
749
750 parser.add_argument(
751 "-n",
752 "--num_threads",
753 required=False,
754 nargs="+",
755 type=int,
756 default=[0],
757 help="Threads to use",
758 )
759
760 parser.add_argument(
761 "--force_num_layers",
762 required=False,
763 type=int,
764 default=None,
765 help="Manually set the model's layer number",
766 )
767
768 parser.add_argument(
769 "--enable_arm64_bfloat16_fastmath_mlas_gemm",
770 required=False,
771 action="store_true",
772 help="Enable bfloat16 mlas gemm kernels on aarch64. Supported only for CPU EP ",
773 )
774 parser.set_defaults(enable_arm64_bfloat16_fastmath_mlas_gemm=False)
775
776 FusionOptions.add_arguments(parser)
777
778 args = parser.parse_args()
779 return args
780
781
782def main():
783 args = parse_arguments()
784
785 setup_logger(args.verbose)
786
787 if args.precision == Precision.FLOAT16 and not args.use_gpu:
788 logger.error("fp16 is for GPU only")
789 return
790
791 if args.precision == Precision.INT8 and args.use_gpu and args.provider not in ["migraphx"]:
792 logger.error("int8 is for CPU only")
793 return
794
795 if len(args.models) == 1 and MODELS[args.models[0]][3] in ["vit", "swim"]:
796 args.sequence_lengths = [""]
797
798 args.num_threads = sorted({cpu_count if x <= 0 else x for x in args.num_threads})
799
800 logger.info(f"Arguments: {args}")
801
802 if not os.path.exists(args.cache_dir):
803 try:
804 os.mkdir(args.cache_dir)
805 except OSError:
806 logger.error("Creation of the directory %s failed", args.cache_dir)
807
808 enable_torch = "torch" in args.engines
809 enable_torch2 = "torch2" in args.engines
810 enable_torchscript = "torchscript" in args.engines
811 enable_onnxruntime = "onnxruntime" in args.engines
812 enable_tensorflow = "tensorflow" in args.engines
813
814 if enable_torch2 and version.parse(torch.__version__) < version.parse("2.0.0"):
815 logger.error(f"PyTorch version must be >=2.0.0 and you are using {torch.__version__}")
816 return
817
818 config_modifier = ConfigModifier(args.force_num_layers)
819
820 results = []
821
822 for num_threads in args.num_threads:
823 torch.set_num_threads(num_threads)
824 logger.debug(torch.__config__.parallel_info())
825 if enable_torch or enable_torch2 or enable_torchscript:
826 if args.input_counts != [1]:
827 logger.warning("--input_counts is not implemented for torch or torchscript engine.")
828
829 if enable_torchscript:
830 results += run_pytorch(
831 args.use_gpu,
832 args.models,
833 args.model_class,
834 config_modifier,
835 args.precision,
836 num_threads,
837 args.batch_sizes,
838 args.sequence_lengths,
839 args.test_times,
840 True,
841 False,
842 args.cache_dir,
843 args.verbose,
844 )
845
846 if enable_torch:
847 results += run_pytorch(
848 args.use_gpu,
849 args.models,
850 args.model_class,
851 config_modifier,
852 args.precision,
853 num_threads,
854 args.batch_sizes,
855 args.sequence_lengths,
856 args.test_times,
857 False,
858 False,
859 args.cache_dir,
860 args.verbose,
861 )
862
863 if enable_torch2:
864 results += run_pytorch(
865 args.use_gpu,
866 args.models,
867 args.model_class,
868 config_modifier,
869 args.precision,
870 num_threads,
871 args.batch_sizes,
872 args.sequence_lengths,
873 args.test_times,
874 False,
875 True,
876 args.cache_dir,
877 args.verbose,
878 )
879
880 if enable_tensorflow:
881 results += run_tensorflow(
882 args.use_gpu,
883 args.models,
884 args.model_class,
885 config_modifier,
886 args.precision,
887 num_threads,
888 args.batch_sizes,
889 args.sequence_lengths,
890 args.test_times,
891 args.cache_dir,
892 args.verbose,
893 )
894
895 model_fusion_statistics = {}
896 if enable_onnxruntime:
897 try:
898 use_raw_attention_mask = not args.use_mask_index
899 results += run_onnxruntime(
900 args.use_gpu,
901 args.provider,
902 args.models,
903 args.model_class,
904 config_modifier,
905 args.precision,
906 num_threads,
907 args.batch_sizes,
908 args.sequence_lengths,
909 args.test_times,
910 args.input_counts,
911 args.optimizer_info,
912 args.validate_onnx,
913 args.cache_dir,
914 args.onnx_dir,
915 args.verbose,
916 args.overwrite,
917 args.disable_ort_io_binding,
918 use_raw_attention_mask,
919 model_fusion_statistics,
920 args.model_source,
921 args.enable_arm64_bfloat16_fastmath_mlas_gemm,
922 args,
923 )
924 except Exception:
925 logger.exception("Exception")
926
927 time_stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
928 if model_fusion_statistics:
929 csv_filename = args.fusion_csv or f"benchmark_fusion_{time_stamp}.csv"
930 output_fusion_statistics(model_fusion_statistics, csv_filename)
931
932 if len(results) == 0:
933 if args.batch_sizes != [0]:
934 logger.warning("No any result available.")
935 return
936
937 csv_filename = args.detail_csv or f"benchmark_detail_{time_stamp}.csv"
938 output_details(results, csv_filename)
939
940 csv_filename = args.result_csv or f"benchmark_summary_{time_stamp}.csv"
941 output_summary(results, csv_filename, args)
942
943
944if __name__ == "__main__":
945 main()
946