Converting from Pytorch/Safetensors to ONNX®

Given the advantages described in Onward With ONNX® we’ve taken the opinion that if it runs on ONNX that’s the way we want to go.  So while ONNX has a large model zoo we’ve had to convert a few models by hand.  Many models like YOLOX provide tools that make this a single command:

python tools/export_onnx.py -f yolox_layouts.py --output-name yolox_m_layout.onnx -n yolox-m -c YOLOX_outputs/yolox_layout/best_ckpt.pth

If the model exists in HuggingFace we’ve seen good results using the Optimum Conversion tool.

This can also be used easily with a few short python commands.  The following example uses the microsoft/table-transformer-structure-recognition-v1.1-all

Steps:

1. Create a directory called model

mkdir model
cd model

2. Download the model from huggingface into the model directory

• Navigate to the model’s files tab
• Select the model.safetensors file
• Also download the config.json

3. Create and activate a Python virtual environment back in your root directory

cd ..
python3 -m venv .env
source .env/bin/activate

4. Install the optimum cli tool

pip install optimum

5. Since we’ll be exporting to onnx, also ensure onnx and onnxruntime is installed

pip install onnx onnxruntime

6. Run the conversion cli to put the exported onnx model in the model_onnx directory.  If the task can not be inferred, it may need to be specified. 

optimum-cli export onnx --task object-detection --model model model_onnx/

7. If successful you should see something similar to the following.  The outputs will be different depending on the model

Validating ONNX model model_onnx/model.onnx...
	-[✓] ONNX model output names match reference model (logits, pred_boxes)
	- Validating ONNX Model output "logits":
		-[✓] (2, 125, 7) matches (2, 125, 7)
		-[✓] all values close (atol: 1e-05)
	- Validating ONNX Model output "pred_boxes":
		-[✓] (2, 125, 4) matches (2, 125, 4)
		-[✓] all values close (atol: 1e-05)

The ONNX export succeeded and the exported model was saved at: model_onnx

(.env) (base) bpaulin@bobs-mbp model_onnx % ls -l
total 226368
-rw-r--r--  1 bpaulin  staff      76783 Jun  4 11:48 config.json
-rw-r--r--  1 bpaulin  staff  115819080 Jun  4 11:48 model.onnx

Congratulations, you’ve just exported your first model to ONNX!  You can now use this model with the ONNX Runtime.  Stay tuned as the true power of using ONNX is unlocked when you can also convert all the pre and post data processing steps to smaller dependency trees. 

At Datavolo we’re on a journey to empowering the 10x Data Engineer and sharing knowledge with you along the way.