def parse_arguments():
parser = argparse.ArgumentParser()
parser.add_argument("-m",
"--models",
required=False,
nargs="+",
type=str,
default=["bert-base-cased", "roberta-base", "gpt2"],
choices=list(MODELS.keys()),
help="Pre-trained models in the list: " +
", ".join(MODELS.keys()))
parser.add_argument("--model_source",
required=False,
nargs=1,
type=str,
default='pt',
choices=['pt', 'tf'],
help="Export onnx from pt or tf")
parser.add_argument('--model_class',
required=False,
type=str,
default=None,
choices=list(MODEL_CLASSES),
help='Model type selected in the list: ' +
', '.join(MODEL_CLASSES))
parser.add_argument(
"-e",
"--engines",
required=False,
nargs="+",
type=str,
default=['onnxruntime'],
choices=['onnxruntime', 'torch', 'torchscript', 'tensorflow'],
help="Engines to benchmark")
parser.add_argument("-c",
"--cache_dir",
required=False,
type=str,
default=os.path.join('.', 'cache_models'),
help="Directory to cache pre-trained models")
parser.add_argument("--onnx_dir",
required=False,
type=str,
default=os.path.join('.', 'onnx_models'),
help="Directory to store onnx models")
parser.add_argument("-g",
"--use_gpu",
required=False,
action="store_true",
help="Run on cuda device")
parser.add_argument(
"-p",
"--precision",
type=Precision,
default=Precision.FLOAT32,
choices=list(Precision),
help=
"Precision of model to run. fp32 for full precision, fp16 for half precision, and int8 for quantization"
)
parser.add_argument("--verbose",
required=False,
action="store_true",
help="Print more information")
parser.add_argument("--overwrite",
required=False,
action="store_true",
help="Overwrite existing models")
parser.add_argument("-o",
"--optimize_onnx",
required=False,
action="store_true",
help="Use optimizer.py to optimize onnx model")
parser.add_argument("-v",
"--validate_onnx",
required=False,
action="store_true",
help="Validate ONNX model")
parser.add_argument(
"-f",
"--fusion_csv",
required=False,
default=None,
help="CSV file for saving summary results of graph optimization.")
parser.add_argument("-d",
"--detail_csv",
required=False,
default=None,
help="CSV file for saving detail results.")
parser.add_argument("-r",
"--result_csv",
required=False,
default=None,
help="CSV file for saving summary results.")
parser.add_argument(
"-i",
"--input_counts",
required=False,
nargs="+",
default=[1],
type=int,
choices=[1, 2, 3],
help=
"Number of ONNX model inputs. Please use 1 for fair comparison with Torch or TorchScript."
)
parser.add_argument(
"-t",
"--test_times",
required=False,
default=100,
type=int,
help="Number of repeat times to get average inference latency.")
parser.add_argument("-b",
"--batch_sizes",
nargs="+",
type=int,
default=[1])
parser.add_argument("-s",
"--sequence_lengths",
nargs="+",
type=int,
default=[4, 8, 16, 32, 64, 128, 256])
parser.add_argument(
'--disable_ort_io_binding',
required=False,
action='store_true',
help='Disable running ONNX Runtime with binded inputs and outputs. ')
parser.set_defaults(disable_ort_io_binding=False)
parser.add_argument("-n",
"--num_threads",
required=False,
nargs="+",
type=int,
default=[0],
help="Threads to use")
args = parser.parse_args()
return args
def parse_arguments():
parser = argparse.ArgumentParser()
parser.add_argument(
"-m",
"--models",
required=False,
nargs="+",
type=str,
default=["bert-base-cased", "roberta-base", "gpt2"],
choices=list(MODELS.keys()),
help="Pre-trained models in the list: " + ", ".join(MODELS.keys()),
)
parser.add_argument(
"--model_source",
required=False,
nargs=1,
type=str,
default="pt",
choices=["pt", "tf"],
help="Export onnx from pt or tf",
)
parser.add_argument(
"--model_class",
required=False,
type=str,
default=None,
choices=list(MODEL_CLASSES),
help="Model type selected in the list: " + ", ".join(MODEL_CLASSES),
)
parser.add_argument(
"-e",
"--engines",
required=False,
nargs="+",
type=str,
default=["onnxruntime"],
choices=["onnxruntime", "torch", "torchscript", "tensorflow"],
help="Engines to benchmark",
)
parser.add_argument(
"-c",
"--cache_dir",
required=False,
type=str,
default=os.path.join(".", "cache_models"),
help="Directory to cache pre-trained models",
)
parser.add_argument(
"--onnx_dir",
required=False,
type=str,
default=os.path.join(".", "onnx_models"),
help="Directory to store onnx models",
)
parser.add_argument("-g",
"--use_gpu",
required=False,
action="store_true",
help="Run on gpu device")
parser.add_argument(
"--provider",
required=False,
type=str,
default=None,
help="Execution provider to use",
)
parser.add_argument(
"-p",
"--precision",
type=Precision,
default=Precision.FLOAT32,
choices=list(Precision),
help=
"Precision of model to run. fp32 for full precision, fp16 for half precision, and int8 for quantization",
)
parser.add_argument("--verbose",
required=False,
action="store_true",
help="Print more information")
parser.add_argument(
"--overwrite",
required=False,
action="store_true",
help="Overwrite existing models",
)
parser.add_argument(
"-o",
"--optimizer_info",
type=OptimizerInfo,
default=OptimizerInfo.BYSCRIPT,
choices=list(OptimizerInfo),
help=
"Optimizer info: Use optimizer.py to optimize onnx model as default. Can also choose from by_ort and no_opt",
)
parser.add_argument(
"-v",
"--validate_onnx",
required=False,
action="store_true",
help="Validate ONNX model",
)
parser.add_argument(
"-f",
"--fusion_csv",
required=False,
default=None,
help="CSV file for saving summary results of graph optimization.",
)
parser.add_argument(
"-d",
"--detail_csv",
required=False,
default=None,
help="CSV file for saving detail results.",
)
parser.add_argument(
"-r",
"--result_csv",
required=False,
default=None,
help="CSV file for saving summary results.",
)
parser.add_argument(
"-i",
"--input_counts",
required=False,
nargs="+",
default=[1],
type=int,
choices=[1, 2, 3],
help=
"Number of ONNX model inputs. Please use 1 for fair comparison with Torch or TorchScript.",
)
parser.add_argument(
"-t",
"--test_times",
required=False,
default=100,
type=int,
help="Number of repeat times to get average inference latency.",
)
parser.add_argument("-b",
"--batch_sizes",
nargs="+",
type=int,
default=[1])
parser.add_argument(
"-s",
"--sequence_lengths",
nargs="+",
type=int,
default=[4, 8, 16, 32, 64, 128, 256],
)
parser.add_argument(
"--disable_ort_io_binding",
required=False,
action="store_true",
help="Disable running ONNX Runtime with binded inputs and outputs. ",
)
parser.set_defaults(disable_ort_io_binding=False)
parser.add_argument(
"-n",
"--num_threads",
required=False,
nargs="+",
type=int,
default=[0],
help="Threads to use",
)
parser.add_argument(
"--force_num_layers",
required=False,
type=int,
default=None,
help="Manually set the model's layer number",
)
FusionOptions.add_arguments(parser)
args = parser.parse_args()
return args