feats: add TensorRT related. ⬆️

CMakeLists.txt

@@ -27,6 +27,12 @@ if(ENABLE_OPENVINO)
     add_definitions(-DUSE_OPENVINO)
 endif(ENABLE_OPENVINO)
 
+option(ENABLE_TENSORRT OFF)
+if(ENABLE_TENSORRT)
+    message("build project with TensorRT")
+    add_definitions(-DUSE_TENSORRT)
+endif(ENABLE_TENSORRT)
+
 include_directories("${ONNXRUNTIME_ROOTDIR}/include"
                 "${CMAKE_SOURCE_DIR}/src"
                 "${CMAKE_SOURCE_DIR}/src/pbrain-Z2I/toml11"

src/common.h

@@ -29,6 +29,7 @@ SOFTWARE.
 // #define SMALL_BOARD_MODE
 // #define USE_CUDA
 // #define USE_OPENVINO
+// #define USE_TENSORRT
 
 #define CHANNEL_SIZE 3
 

src/onnx.cpp

@@ -95,6 +95,14 @@ NeuralNetwork::NeuralNetwork(const std::string model_path, const unsigned int ba
   CheckStatus(g_ort, g_ort->SessionOptionsAppendExecutionProvider_OpenVINO(session_options, &OpenVINO_Options));
 #endif
 
+#ifdef USE_TENSORRT
+  const OrtApiBase* ptr_api_base = OrtGetApiBase();
+  const OrtApi* g_ort = ptr_api_base->GetApi(ORT_API_VERSION);
+  OrtTensorRTProviderOptions TensorRT_Options;
+  TensorRT_Options.device_id = 0;
+  CheckStatus(g_ort, g_ort->SessionOptionsAppendExecutionProvider_TensorRT(session_options, &TensorRT_Options));
+#endif
+
 #ifdef _WIN32
   // std::wstring widestr = std::wstring(model_path.begin(), model_path.end());
   std::wstring wstr(model_path.length(), L' ');