onnxruntime/samples/python/pytorch_transformer
2020-10-12 17:53:09 -07:00
..
ort_train.py Update PyTorch TransformerModel sample (#5275) 2020-09-24 16:28:07 -07:00
ort_utils.py Add new PytTrch front-end (#4815) 2020-08-17 09:45:25 -07:00
pt_model.py Enable example transformer test with dynamic size inputs (#4888) 2020-08-24 14:31:08 -07:00
pt_train.py Update PyTorch TransformerModel sample (#5275) 2020-09-24 16:28:07 -07:00
README.md Fix commands in README.md. (#5459) 2020-10-12 17:53:09 -07:00
utils.py Add new PytTrch front-end (#4815) 2020-08-17 09:45:25 -07:00

TransformerModel example

This example was adapted from Pytorch's Sequence-to-Sequence Modeling with nn.Transformer and TorchText tutorial

Requirements

  • PyTorch 1.6+
  • TorchText 0.6+
  • ONNX Runtime 1.5+

Running PyTorch version

python pt_train.py

Running ONNX Runtime version

python ort_train.py

Optional arguments

Argument Description Default
--batch-size input batch size for training 20
--test-batch-size input batch size for testing 20
--epochs number of epochs to train 2
--lr learning rate 0.001
--no-cuda disables CUDA training False
--seed random seed 1
--log-interval how many batches to wait before logging training status 200