Skip to content

Commit 22675e9

Browse files
JuliaBareevaJuliaBareeva
committed
add code for the article How to convert a model from PyTorch to TensorRT and speed up inference
1 parent dfc798e commit 22675e9

File tree

8 files changed

+1424
-0
lines changed

8 files changed

+1424
-0
lines changed

‎PyTorch-ONNX-TensorRT/CMakeLists.txt‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
cmake_minimum_required(VERSION 3.10)
2+
set(CMAKE_CXX_STANDARD 14)
3+
set(CMAKE_CXX_STANDARD_REQUIRED TRUE)
4+
set(CMAKE_MODULE_PATH "${CMAKE_SOURCE_DIR}" ${CMAKE_MODULE_PATH})
5+
6+
project(TensorRTSample)
7+
8+
find_package(Threads)
9+
find_package(CUDA REQUIRED)
10+
find_package(OpenCV REQUIRED)
11+
find_package(TensorRT REQUIRED)
12+
13+
# TARGETS --------------------------------------------------------------------------------------------------------------
14+
add_executable(trt_sample trt_sample.cpp)
15+
target_include_directories(trt_sample PUBLIC ${OpenCV_INCLUDE_DIRS} ${CUDA_INCLUDE_DIRS} ${TensorRT_INCLUDE_DIRS})
16+
target_link_libraries(trt_sample PUBLIC ${OpenCV_LIBS} ${CUDA_LIBRARIES} ${CMAKE_THREAD_LIBS_INIT} ${TensorRT_LIBRARIES})

‎PyTorch-ONNX-TensorRT/FindTensorRT.cmake‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
# source:
2+
# https://github.com/NVIDIA/tensorrt-laboratory/blob/master/cmake/FindTensorRT.cmake
3+
4+
# This module defines the following variables:
5+
#
6+
# ::
7+
#
8+
# TensorRT_INCLUDE_DIRS
9+
# TensorRT_LIBRARIES
10+
# TensorRT_FOUND
11+
#
12+
# ::
13+
#
14+
# TensorRT_VERSION_STRING - version (x.y.z)
15+
# TensorRT_VERSION_MAJOR - major version (x)
16+
# TensorRT_VERSION_MINOR - minor version (y)
17+
# TensorRT_VERSION_PATCH - patch version (z)
18+
#
19+
# Hints
20+
# ^^^^^
21+
# A user may set ``TensorRT_DIR`` to an installation root to tell this module where to look.
22+
#
23+
set(_TensorRT_SEARCHES)
24+
25+
if(TensorRT_DIR)
26+
set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_DIR} NO_DEFAULT_PATH)
27+
list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT)
28+
endif()
29+
30+
# appends some common paths
31+
set(_TensorRT_SEARCH_NORMAL
32+
PATHS "/usr"
33+
)
34+
list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL)
35+
36+
# Include dir
37+
foreach(search ${_TensorRT_SEARCHES})
38+
find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include)
39+
endforeach()
40+
41+
if(NOT TensorRT_LIBRARY)
42+
foreach(search ${_TensorRT_SEARCHES})
43+
find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib)
44+
endforeach()
45+
endif()
46+
47+
if(NOT TensorRT_PARSERS_LIBRARY)
48+
foreach(search ${_TensorRT_SEARCHES})
49+
find_library(TensorRT_NVPARSERS_LIBRARY NAMES nvparsers ${${search}} PATH_SUFFIXES lib)
50+
endforeach()
51+
endif()
52+
53+
if(NOT TensorRT_NVONNXPARSER_LIBRARY)
54+
foreach(search ${_TensorRT_SEARCHES})
55+
find_library(TensorRT_NVONNXPARSER_LIBRARY NAMES nvonnxparser ${${search}} PATH_SUFFIXES lib)
56+
endforeach()
57+
endif()
58+
59+
mark_as_advanced(TensorRT_INCLUDE_DIR)
60+
61+
if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h")
62+
file(STRINGS "${TensorRT_INCLUDE_DIR}/NvInfer.h" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$")
63+
file(STRINGS "${TensorRT_INCLUDE_DIR}/NvInfer.h" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$")
64+
file(STRINGS "${TensorRT_INCLUDE_DIR}/NvInfer.h" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$")
65+
66+
string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}")
67+
string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}")
68+
string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}")
69+
set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}")
70+
endif()
71+
72+
include(FindPackageHandleStandardArgs)
73+
FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING)
74+
75+
if(TensorRT_FOUND)
76+
set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR})
77+
78+
if(NOT TensorRT_LIBRARIES)
79+
set(TensorRT_LIBRARIES ${TensorRT_LIBRARY} ${TensorRT_NVONNXPARSER_LIBRARY} ${TensorRT_NVPARSERS_LIBRARY})
80+
endif()
81+
82+
if(NOT TARGET TensorRT::TensorRT)
83+
add_library(TensorRT::TensorRT UNKNOWN IMPORTED)
84+
set_target_properties(TensorRT::TensorRT PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}")
85+
set_property(TARGET TensorRT::TensorRT APPEND PROPERTY IMPORTED_LOCATION "${TensorRT_LIBRARY}")
86+
endif()
87+
endif()

‎PyTorch-ONNX-TensorRT/README.md‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
# How to convert a model from PyTorch to TensorRT and speed up inference
2+
The blog post is here: https://www.learnopencv.com/how-to-convert-a-model-from-pytorch-to-tensorrt-and-speed-up-inference/
3+
4+
To run Python part:
5+
```shell script
6+
python3 -m pip install -r requirements.txt
7+
python3 pytorch_model.py
8+
```
9+
10+
To run C++ part:
11+
```shell script
12+
mkdir build
13+
cd build
14+
cmake -DOpenCV_DIR=[path-to-opencv-build] -DTensorRT_DIR=[path-to-tensorrt] ..
15+
make -j8
16+
trt_sample[.exe] resnet50.onnx turkish_coffee.jpg
17+
```
18+
19+
# AI Courses by OpenCV
20+
21+
Want to become an expert in AI? [AI Courses by OpenCV](https://opencv.org/courses/) is a great place to start.
22+
23+
<a href="https://opencv.org/courses/">
24+
<p align="center">
25+
<img src="https://www.learnopencv.com/wp-content/uploads/2020/04/AI-Courses-By-OpenCV-Github.png">
26+
</p>
27+
</a>

0 commit comments

Comments
 (0)