Note: In this tutorial, we assume nnfusion cli has been installed as Build Guide.
NNFusion leverages ONNX to support PyTorch. So this tutorial focuses on how to freeze an ONNX model from PyTorch source code. You could get NNFusion supported ONNX ops here.
Please refer to PyTorch onnx section to convert a PyTorch model to ONNX format, currently it already supports a great majority of deeplearning workloads.
On PyTorch onnx_exporter, we build a simple wrapper called pt_freezer, it wraps PyTorch onnx_exporter with control flow and op availability(not implemented yet) check. We provide a well self-explanatory VGG example for this tool:
# step0: install prerequisites
apt update && sudo apt install python3-pip
pip3 install onnx torch torchvision
# step1: freeze vgg16 model
python3 vgg16_model.py
Of course, you could freeze ONNX model from thirdparties, like huggingface transformer, which supports exporting to ONNX format.
model | nnf codegen flags | download link |
---|---|---|
VGG16 | -f onnx | vgg16.onnx |
BERT_base | -f onnx -p 'batch:3;sequence:512' | pt-bert-base-cased.onnx |