Skip to content

Instantly share code, notes, and snippets.

@Emekaborisama

Emekaborisama/convert_transformer_to_onnx.py

Created August 30, 2022 12:39
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save Emekaborisama/515890e81d6dcf7a96d2906050c0ed11 to your computer and use it in GitHub Desktop.
Save Emekaborisama/515890e81d6dcf7a96d2906050c0ed11 to your computer and use it in GitHub Desktop.
Download ZIP
convert transformers model to onnx using pytorch
Raw
convert_transformer_to_onnx.py
torch.onnx.export(
model,
tuple(encoded_input.values()),
f="torch-model.onnx",
input_names=['input_ids', 'attention_mask','token_type_ids'],
output_names=['logits'],
dynamic_axes={'input_ids': {0: 'batch_size', 1: 'sequence'},
'attention_mask': {0: 'batch_size', 1: 'sequence'},
'token_type_ids': {0: 'batch_size', 1: 'sequence'},
'logits': {0: 'batch_size', 1: 'sequence'}},
do_constant_folding=True,
opset_version=13,
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment