def main():
import benchmark_helper
args = docopt.docopt(__doc__)
kwargs = {
'model_name': args['<model_name>'],
'seq_len': int(args['--seq_len']),
'batch_size': int(args['--batch_size']),
'n': int(args['-n']),
}
if args['--framework'] == 'turbo-transformers':
benchmark_turbo_transformers(**kwargs)
elif args['--framework'] == 'torch':
benchmark_torch(**kwargs)
elif args['--framework'] == 'onnxruntime':
benchmark_helper.onnxruntime_benchmark_creator('GPU')(**kwargs)
else:
raise RuntimeError(f"Not supportted framework {args['--framework']}")
def main():
args = docopt.docopt(__doc__)
kwargs = {
'model': args['<model>'],
'seq_len': int(args['--seq_len']),
'batch_size': int(args['--batch_size']),
'n': int(args['-n']),
'num_threads': int(args['--num_threads'])
}
if args['--framework'] == 'turbo-transformers':
benchmark_turbo_transformers(**kwargs)
elif args['--framework'] == 'torch':
benchmark_torch(**kwargs)
elif args['--framework'] == 'torch_jit':
benchmark_torch_jit(**kwargs)
elif args['--framework'] == 'onnxruntime-cpu':
import benchmark_helper
benchmark_helper.onnxruntime_benchmark_creator('CPU')(**kwargs)
elif args['--framework'] == 'onnxruntime-mkldnn':
import benchmark_helper
benchmark_helper.onnxruntime_benchmark_creator('MKL-DNN')(**kwargs)
else:
raise RuntimeError(f"Not supportted framework {args['--framework']}")