https://github.com/dataxujing/quick-demo-from-pytorch-to-tensort-cpp
A simple and quick example shows how to convert the model of Pytoch to ONNX, and then deploy it with tensorrt in C++
https://github.com/dataxujing/quick-demo-from-pytorch-to-tensort-cpp
Last synced: about 1 year ago
JSON representation
A simple and quick example shows how to convert the model of Pytoch to ONNX, and then deploy it with tensorrt in C++
- Host: GitHub
- URL: https://github.com/dataxujing/quick-demo-from-pytorch-to-tensort-cpp
- Owner: DataXujing
- Created: 2021-06-04T14:21:04.000Z (about 5 years ago)
- Default Branch: main
- Last Pushed: 2021-02-16T07:33:00.000Z (over 5 years ago)
- Last Synced: 2025-02-09T09:29:08.109Z (over 1 year ago)
- Homepage:
- Size: 88.9 KB
- Stars: 4
- Watchers: 1
- Forks: 1
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
# Quick-Demo-from-Pytorch-to-Tensort-cpp
A simple and quick example shows how to convert a pytorch's model to ONNX, and then deploy it with tensorrt using c++.
## Pipline
- **step 1.** Using Pytorch to build a simple neural network and then export to ONNX file `test.onnx`
```bash
git clone https://github.com/hova88/Quick-Demo-from-Pytorch-to-Tensort-cpp.git
cd Quick-Demo-from-Pytorch-to-Tensort-cpp
python pytorch_to_onnx.py
```
The screen will shown as this:
```
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 10, 24, 24] 260
Conv2d-2 [-1, 20, 8, 8] 5,020
Dropout2d-3 [-1, 20, 8, 8] 0
Linear-4 [-1, 50] 16,050
Linear-5 [-1, 10] 510
================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.06
Params size (MB): 0.08
Estimated Total Size (MB): 0.15
----------------------------------------------------------------
Dummy output:
tensor([[-2.2734, -2.2539, -2.2107, -2.3713, -2.3443, -2.4683, -2.3709, -2.2777,
-2.2169, -2.2676]], device='cuda:0')
```
- **step 2.** Initialize tensorrt engine by ONNX file, and then DO INFERENCE.
```bash
mkdir build && cd build && cmake ..
make -j6
./onnxTotrt
```
The screen will shown as this:
```bash
----------------------------------------------------------------
Input filename: ../test.onnx
ONNX IR version: 0.0.4
Opset version: 10
Producer name: pytorch
Producer version: 1.3
Domain:
Model version: 0
Doc string:
----------------------------------------------------------------
WARNING: [TRT]/home/hova/onnx-tensorrt/onnx2trt_utils.cpp:220: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
loaded trt model , do inference
Check output:
-2.20698 -2.24908 -2.4583 -2.27005 -2.30133 -2.43457 -2.33783 -2.21271 -2.41685 -2.1832
```