El modelo de aprendizaje profundo PyTorch se entrena y se transfiere a ONNX y TensorRT para su implementación.
En términos generales, el modelo de aprendizaje profundo se entrena en función de Python. Cuando se implementa el modelo, generalmente debe implementarse en función del código C ++. Este artículo toma a LeNet como ejemplo para presentar todo el proceso de entrenamiento de un modelo de aprendizaje profundo desde el entrenamiento basado en PyTorch hasta la implementación de TensorRT. Para obtener el código detallado, consulte: https://github.com/linghu8812/tensorrt_inference/tree/master/ lenet .
1. Modelo de formación
python3 train.pyUn modelo que reconoce dígitos escritos a mano se puede entrenar mediante comandos , y la precisión de reconocimiento del modelo es de aproximadamente el 99%.
2. Convierta el modelo PyTorch en el modelo ONNX
python3 export_onnx.pyEl modelo PyTorch se puede convertir a un modelo ONNX mediante comandos . El marco TensorRT no puede analizar directamente el modelo PyTorch, por lo que primero convierta el modelo PyTorch. En el siguiente código, el tamaño de lote de la conversión es 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
Después de clonar el código, compílelo directamente.
mkdir build && cd build
cmake ..
make -j
4. Ejecute lenet_trt
Después de compilar, ejecútelo directamente ./lenet_trt ../config.yaml ../samples/. sampleEn la carpeta están las imágenes que deben reconocerse. config.yamlDefine la configuración de la operación del modelo, como el archivo onnx que se va a convertir y el nombre del archivo tensorrt convertido, el tamaño del lote durante la inferencia y el ancho, alto y número de canales durante la inferencia del modelo.
lenet:
onnx_file: "../mnist_net.onnx"
engine_file: "../mnist_net.trt"
BATCH_SIZE: 10
INPUT_CHANNEL: 1
IMAGE_WIDTH: 28
IMAGE_HEIGHT: 28
Los resultados de la operación son los siguientes:
