PyTorch to TensorRT C++ Conversion

This repository provides a step-by-step guide and code for converting a PyTorch model to a TensorRT engine and running it in C++. The process includes exporting the PyTorch model to the ONNX format, generating a TensorRT engine, and performing inference. In this example, we demonstrate the conversion using a simple convolutional neural network for recognizing MNIST handwritten digits.

Table of Contents

• Requirements
• Getting Started
• Conversion Steps
• Running the Inference
• References

Requirements

Before you begin, ensure that you have the following prerequisites installed:

• Python 3.x
• PyTorch
• ONNX
• CUDA (for GPU support)
• TensorRT

Getting Started

1. Clone this repository:

   git clone https://github.com/albertnew2012/pytorch-onnx-tensorRT.git
   cd pytorch-onnx-tensorRT
   

2. Create a simple convolutional model in PyTorch and export it to ONNX format

   python3 ./scripts/simpleCNN_gpu.py

Conversion Steps

The conversion process involves the following steps:

1. Export PyTorch Model to ONNX format: This step converts your PyTorch model to the ONNX format, which is compatible with TensorRT.

   /usr/src/tensorrt/bin/trtexec --onnx=./simple_cnn.onnx --saveEngine=./simple_cnn.engine 

2. Build the C++ Inference Code: Set up the C++ environment for running the TensorRT inference. Build the C++ code using CMake.

3. Generate TensorRT Engine: Use the C++ code to generate a TensorRT engine from the ONNX model.

Running the Inference

1. Navigate to the src directory

   cd src

2. Create a build directory and build the C++ code

   mkdir build && cd build
   cmake ..
   make

3. Run the TensorRT inference

   ./tensorrt_inference

This program will load the TensorRT engine, perform inference on the MNIST dataset, and display the results.

References

• PyTorch
• ONNX
• NVIDIA CUDA Toolkit
• NVIDIA TensorRT