Export the ONNX model files to disk#
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.
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
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.
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)
export_onnx_fileis called, the forecaster will automatically build an ONNXRuntime session with default settings. So you can directly call this method without calling
predict_with_onnxfirst. But when you want to export quantized onnx model files, you should quantize the forecaster by calling
If you just need to export fp32 onnx files, you could specify
dirnameonly and set
The files exported will be saved at
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.