saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Force return_tuple=True to handle transformers breaking change of output format. (#4599)

Browse source
This commit is contained in:
Tianlei Wu 2020-07-23 11:35:41 -07:00 committed by GitHub
parent ea87c0d028
commit ace41b8064
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 8 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/python/tools/transformers/benchmark_gpt2.py
View file

@ -126,6 +126,8 @@ def main():
model_class = MODEL_CLASSES[args.model_class][0]
config = AutoConfig.from_pretrained(args.model_name, torchscript=args.torchscript, cache_dir=cache_dir)
if hasattr(config, 'return_tuple'):
config.return_tuple = True
model = model_class.from_pretrained(args.model_name, config=config, cache_dir=cache_dir)
# This scirpt does not support float16 for PyTorch.

5
onnxruntime/python/tools/transformers/convert_to_onnx.py
View file

@ -111,7 +111,10 @@ def main():
assert not args.use_gpu, "quantization only supports CPU"
model_class = MODEL_CLASSES[args.model_class][0]
model = model_class.from_pretrained(args.model_name_or_path, cache_dir=cache_dir)
config = AutoConfig.from_pretrained(args.model_name_or_path, cache_dir=cache_dir)
if hasattr(config, 'return_tuple'):
config.return_tuple = True
model = model_class.from_pretrained(args.model_name_or_path, config=config, cache_dir=cache_dir)
device = torch.device("cuda:0" if args.use_gpu else "cpu")
model.eval().to(device)

2
onnxruntime/python/tools/transformers/onnx_exporter.py
View file

@ -167,6 +167,8 @@ def export_onnx_model(model_name, opset_version, use_external_data_format, model
use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite,
model_fusion_statistics):
config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir)
if hasattr(config, 'return_tuple'):
config.return_tuple = True
model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir)
model.cpu()