Export the ONNX model files to disk#

Introduction#

When a forecaster is accelerated by ONNXRuntime, we can save the ONNX model files to disk by calling export_onnx_file. In this guidance, we demonstrate how to export the ONNX model files to disk in detail.

We will take TCNForecaster and nyc_taxi dataset as an example in this guide.

Setup#

Before we begin, we need to install Chronos if it isn’t already available, we choose to use pytorch as deep learning backend.

[ ]:

!pip install --pre --upgrade bigdl-chronos[pytorch]
# install ONNXRuntime
!pip install onnx
!pip install onnxruntime
# uninstall torchtext to avoid version conflict
!pip uninstall -y torchtext
# for quantization
!pip install neural-compressor
!pip install onnxruntime-extensions

📝Note

Although Chronos supports inferencing on a cluster, the method to export model files can only be used when forecaster is a non-distributed version.

Only pytorch backend deep learning forecasters support onnxruntime acceleration.

Forecaster preparation#

Before the exporting process, a forecaster should be created and trained. The training process is introduced in the previous guidance Train forcaster on single node in detail, therefore we directly create and train a TCNForecaster based on the nyc taxi dataset.

Export the ONNX model files#

When a trained forecaster is ready and forecaster is a non-distributed version, we provide with export_onnx_file method to export the ONNX model files to disk. The export_onnx_file method has 2 parameters: dirname is the location to save the ONNX files, and quantized_dirname is the location to save the quantized ONNX files when you have a quantized forecaster.

[ ]:

from pathlib import Path

# get data for training and testing and validating
train_data, test_data, val_data = get_data()
# get a trained forecaster
forecaster = get_trained_forecaster(train_data)

# quantize the forecaster
forecaster.quantize(train_data, val_data=val_data,
                    framework='onnxrt_qlinearops')

# create a directory to save onnx files
dirname = Path("onnx_files")
dirname.mkdir(exist_ok=True)
ckpt_name = dirname / "fp32_onnx"
ckpt_name_q = dirname / "int_onnx"

# export the onnx files
forecaster.export_onnx_file(dirname=ckpt_name, quantized_dirname=ckpt_name_q)

📝Note

When export_onnx_file is called, the forecaster will automatically build an ONNXRuntime session with default settings. So you can directly call this method without calling predict_with_onnx first. But when you want to export quantized onnx model files, you should quantize the forecaster by calling quantize method first.

If you just need to export fp32 onnx files, you could specify dirname only and set quantized_dirname to None: forecaster.export_openvino_file(dirname=ckpt_name, quantized_dirname=None)

The files exported will be saved at onnx_files directory.

There are 2 files in each subdirectory:

nano_model_meta.yml: meta information of the saved model checkpoint
onnx_saved_model.onnx: model checkpoint for general use, describes model structure

You only need to take onnx_saved_model.onnx for futher usage.