The deep learning model PyTorch is trained and transferred to ONNX and TensorRT for deployment
Generally speaking, deep learning models are trained based on python. When the models are deployed, they generally need to be deployed based on C++ code. This article takes LeNet as an example to introduce the whole process of training a deep learning model from PyTorch-based training to TensorRT deployment. For the detailed code, see: https://github.com/linghu8812/tensorrt_inference/tree/master/lenet .
1. Model training
python3 train.pyA model that recognizes handwritten digits can be trained through commands , and the recognition accuracy of the model is about 99%.
2. Convert PyTorch model to ONNX model
python3 export_onnx.pyThe PyTorch model can be converted to an ONNX model through commands . The TensorRT framework cannot directly parse the PyTorch model, so first convert the PyTorch model. In the following code, the batch size of the conversion is 10.
import onnx
import torch
# export from pytorch to onnx
net = torch.load('mnist_net.pt').to('cpu')
image = torch.randn(10, 1, 28, 28)
torch.onnx.export(net, image, 'mnist_net.onnx', input_names=['input'], output_names=['output'])
# check onnx model
onnx_model = onnx.load("mnist_net.onnx") # load onnx model
onnx.checker.check_model(onnx_model)
3. Compile lenet_trt
After cloning the code, just compile it directly.
mkdir build && cd build
cmake ..
make -j
4. Run lenet_trt
After compiling, run it directly ./lenet_trt ../config.yaml ../samples/. sampleIn the folder are the pictures that need to be recognized. config.yamlDefines the configuration of the model operation, such as the onnx file to be converted and the converted tensorrt file name, the batch size during inference, and the width, height and number of channels during model inference.
lenet:
onnx_file: "../mnist_net.onnx"
engine_file: "../mnist_net.trt"
BATCH_SIZE: 10
INPUT_CHANNEL: 1
IMAGE_WIDTH: 28
IMAGE_HEIGHT: 28
The results of the operation are as follows:
