add saveEngine for save modle.engine from onnx parser

Author: zhoucoolqi

.gitignore

@@ -9,3 +9,4 @@ bin/
 *.vcxproj
 externals/
 **/.DS_Store
+save

.vscode/launch.json

@@ -9,7 +9,7 @@
  "type": "cppdbg",
  "request": "launch",
  "program": "/home/zqi/sbx/tensorrt-github-samples/build/sample_onnx_mnist_debug",
- "args": ["-d", "/home/zqi/sbx/data/tensorrt-sample-data/mnist/", "--fp16"],
+ "args": ["-d", "/home/zqi/sbx/data/tensorrt-sample-data/mnist/", "--fp16", "--saveEngine", "/home/zqi/sbx/tensorrt-github-samples/save/mnist.engine"],
  "environment": [{ "name": "config", "value": "Debug" }],
  "MIMode": "gdb",
  "miDebuggerPath": "/usr/bin/gdb", 

CMakeLists.txt

@@ -38,20 +38,20 @@ set(RUNTIME_OUTPUT_DIRECTORY ${TRT_OUT_DIR} CACHE PATH "Output directory for run
 set(LIBRARY_OUTPUT_DIRECTORY ${TRT_OUT_DIR} CACHE PATH "Output directory for library target files")
 set(ARCHIVE_OUTPUT_DIRECTORY ${TRT_OUT_DIR} CACHE PATH "Output directory for archive target files")
 
-set(STATIC_LIB_EXT "a")
+# set(STATIC_LIB_EXT "a")
 
-# Get tensorrt version info
-file(STRINGS "${CMAKE_CURRENT_SOURCE_DIR}/include/NvInferVersion.h" VERSION_STRINGS REGEX "#define NV_TENSORRT_.*")
+# # Get tensorrt version info
+# file(STRINGS "${CMAKE_CURRENT_SOURCE_DIR}/include/NvInferVersion.h" VERSION_STRINGS REGEX "#define NV_TENSORRT_.*")
 
-foreach(TYPE MAJOR MINOR PATCH BUILD)
- string(REGEX MATCH "NV_TENSORRT_${TYPE} [0-9]+" TRT_TYPE_STRING ${VERSION_STRINGS})
- string(REGEX MATCH "[0-9]+" TRT_${TYPE} ${TRT_TYPE_STRING})
-endforeach(TYPE)
+# foreach(TYPE MAJOR MINOR PATCH BUILD)
+#  string(REGEX MATCH "NV_TENSORRT_${TYPE} [0-9]+" TRT_TYPE_STRING ${VERSION_STRINGS})
+#  string(REGEX MATCH "[0-9]+" TRT_${TYPE} ${TRT_TYPE_STRING})
+# endforeach(TYPE)
 
-set(TRT_VERSION "${TRT_MAJOR}.${TRT_MINOR}.${TRT_PATCH}" CACHE STRING "TensorRT project version")
-set(ONNX2TRT_VERSION "${TRT_MAJOR}.${TRT_MINOR}.${TRT_PATCH}" CACHE STRING "ONNX2TRT project version")
-set(TRT_SOVERSION "${TRT_MAJOR}" CACHE STRING "TensorRT library so version")
-message("Building for TensorRT version: ${TRT_VERSION}, library version: ${TRT_SOVERSION}")
+# set(TRT_VERSION "${TRT_MAJOR}.${TRT_MINOR}.${TRT_PATCH}" CACHE STRING "TensorRT project version")
+# set(ONNX2TRT_VERSION "${TRT_MAJOR}.${TRT_MINOR}.${TRT_PATCH}" CACHE STRING "ONNX2TRT project version")
+# set(TRT_SOVERSION "${TRT_MAJOR}" CACHE STRING "TensorRT library so version")
+# message(STATUS "Building for TensorRT version: ${TRT_VERSION}, library version: ${TRT_SOVERSION}")
 
 # Set g++ cmake flag
 if(NOT DEFINED CMAKE_TOOLCHAIN_FILE)
@@ -63,11 +63,12 @@ if(NOT DEFINED CMAKE_TOOLCHAIN_FILE)
 endif()
 
 # Set project info : languages c++ and cuda
-project(TensorRT
- LANGUAGES CXX CUDA
- VERSION ${TRT_VERSION}
- DESCRIPTION "TensorRT is a C++ library that facilitates high-performance inference on NVIDIA GPUs and deep learning accelerators."
- HOMEPAGE_URL "https://github.com/NVIDIA/TensorRT")
+project(TensorRT_zqi)
+# project(TensorRT
+# LANGUAGES CXX CUDA
+# VERSION ${TRT_VERSION}
+# DESCRIPTION "TensorRT is a C++ library that facilitates high-performance inference on NVIDIA GPUs and deep learning accelerators."
+# HOMEPAGE_URL "https://github.com/NVIDIA/TensorRT")
 
 # Set cmake install path ../bin/
 if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
@@ -124,8 +125,8 @@ include_directories(
 )
 
 set(nvinfer_lib_name "nvinfer")
-set(nvinfer_plugin_lib_name "nvinfer_plugin")
-set(nvinfer_vc_plugin_lib_name "nvinfer_vc_plugin")
+# set(nvinfer_plugin_lib_name "nvinfer_plugin")
+# set(nvinfer_vc_plugin_lib_name "nvinfer_vc_plugin")
 set(nvonnxparser_lib_name "nvonnxparser")
 
 find_library_create_target(nvinfer ${nvinfer_lib_name} SHARED ${TRT_LIB_DIR})

CMakeLists.txt.bk.20150105

@@ -0,0 +1,147 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+cmake_minimum_required(VERSION 3.13 FATAL_ERROR)
+include(cmake/modules/set_ifndef.cmake)
+include(cmake/modules/find_library_create_target.cmake)
+
+# Set lib and out path equals current build path, e.g. : tensorrt-github-samples/build
+set_ifndef(TRT_LIB_DIR ${CMAKE_BINARY_DIR})
+set_ifndef(TRT_OUT_DIR ${CMAKE_BINARY_DIR})
+
+if(CMAKE_VERSION VERSION_LESS 3.20)
+ file(TO_CMAKE_PATH "${TRT_LIB_DIR}" TRT_LIB_DIR)
+ file(TO_CMAKE_PATH "${TRT_OUT_DIR}" TRT_OUT_DIR)
+else()
+ cmake_path(SET TRT_LIB_DIR ${TRT_LIB_DIR})
+ cmake_path(SET TRT_OUT_DIR ${TRT_OUT_DIR})
+endif()
+message(STATUS "TRT_LIB_DIR: ${TRT_LIB_DIR}")
+message(STATUS "TRT_OUT_DIR: ${TRT_OUT_DIR}")
+
+# Set compile output paths
+set(RUNTIME_OUTPUT_DIRECTORY ${TRT_OUT_DIR} CACHE PATH "Output directory for runtime target files")
+set(LIBRARY_OUTPUT_DIRECTORY ${TRT_OUT_DIR} CACHE PATH "Output directory for library target files")
+set(ARCHIVE_OUTPUT_DIRECTORY ${TRT_OUT_DIR} CACHE PATH "Output directory for archive target files")
+
+set(STATIC_LIB_EXT "a")
+
+# Get tensorrt version info
+file(STRINGS "${CMAKE_CURRENT_SOURCE_DIR}/include/NvInferVersion.h" VERSION_STRINGS REGEX "#define NV_TENSORRT_.*")
+
+foreach(TYPE MAJOR MINOR PATCH BUILD)
+ string(REGEX MATCH "NV_TENSORRT_${TYPE} [0-9]+" TRT_TYPE_STRING ${VERSION_STRINGS})
+ string(REGEX MATCH "[0-9]+" TRT_${TYPE} ${TRT_TYPE_STRING})
+endforeach(TYPE)
+
+set(TRT_VERSION "${TRT_MAJOR}.${TRT_MINOR}.${TRT_PATCH}" CACHE STRING "TensorRT project version")
+set(ONNX2TRT_VERSION "${TRT_MAJOR}.${TRT_MINOR}.${TRT_PATCH}" CACHE STRING "ONNX2TRT project version")
+set(TRT_SOVERSION "${TRT_MAJOR}" CACHE STRING "TensorRT library so version")
+message(STATUS "Building for TensorRT version: ${TRT_VERSION}, library version: ${TRT_SOVERSION}")
+
+# Set g++ cmake flag
+if(NOT DEFINED CMAKE_TOOLCHAIN_FILE)
+ find_program(CMAKE_CXX_COMPILER NAMES $ENV{CXX} g++) # CMake 的 find_program 行为是 "lazy" 的，只在变量为空或未定义时才会查找并赋值
+ message(STATUS "CMAKE_CXX_COMPILER is: ${CMAKE_CXX_COMPILER}")
+ if(NOT CMAKE_CXX_COMPILER)
+ message(FATAL_ERROR "C++ compiler not found. Please specify one using the CXX environment variable or ensure g++ is installed.")
+ endif()
+endif()
+
+# Set project info : languages c++ and cuda
+project(TensorRT
+ LANGUAGES CXX CUDA
+ VERSION ${TRT_VERSION}
+ DESCRIPTION "TensorRT is a C++ library that facilitates high-performance inference on NVIDIA GPUs and deep learning accelerators."
+ HOMEPAGE_URL "https://github.com/NVIDIA/TensorRT")
+
+# Set cmake install path ../bin/
+if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
+ set(CMAKE_INSTALL_PREFIX ${TRT_LIB_DIR}/../ CACHE PATH "TensorRT installation" FORCE)
+endif(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
+
+option(BUILD_SAMPLES "Build TensorRT samples" ON)
+
+# Set C++ standard version
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_EXTENSIONS OFF)
+
+# Set cmake cxx flags, define: -DBUILD_SYSTEM=cmake_oss
+if(NOT MSVC)
+ set(CMAKE_CXX_FLAGS "-Wno-deprecated-declarations ${CMAKE_CXX_FLAGS} -DBUILD_SYSTEM=cmake_oss")
+else()
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DBUILD_SYSTEM=cmake_oss")
+endif()
+
+# Set plat form info, tbd
+set_ifndef(TRT_PLATFORM_ID "x86_64")
+message(STATUS "Targeting TRT Platform: ${TRT_PLATFORM_ID}")
+
+set(TRT_DEBUG_POSTFIX _debug CACHE STRING "suffix for debug builds")
+
+if (CMAKE_BUILD_TYPE STREQUAL "Debug")
+ message(STATUS "Building in debug mode ${DEBUG_POSTFIX}")
+endif()
+
+# Set dependency: threads, find cuda lib, include cuda CUDA_INCLUDE_DIRS
+set(DEFAULT_CUDA_VERSION 12.2.0)
+set(DEFAULT_CUDNN_VERSION 8.9)
+
+## Dependency Version Resolution
+set_ifndef(CUDA_VERSION ${DEFAULT_CUDA_VERSION})
+message(STATUS "CUDA version set to ${CUDA_VERSION}")
+set_ifndef(CUDNN_VERSION ${DEFAULT_CUDNN_VERSION})
+message(STATUS "cuDNN version set to ${CUDNN_VERSION}")
+
+set(THREADS_PREFER_PTHREAD_FLAG ON)
+find_package(Threads REQUIRED)
+
+message(STATUS "CUDA_INCLUDE_DIRS before find package: ${CUDA_INCLUDE_DIRS}")
+## find_package(CUDA) is broken for cross-compilation. Enable CUDA language instead.
+message(STATUS "CMAKE_PREFIX_PATH: ${CMAKE_PREFIX_PATH}")
+if(NOT DEFINED CMAKE_TOOLCHAIN_FILE)
+ find_package(CUDA ${CUDA_VERSION} REQUIRED EXACT) # 即使精确指定，仍然可能找到本机的12.4的cuda版本
+endif()
+message(STATUS "CUDA_INCLUDE_DIRS after find package: ${CUDA_INCLUDE_DIRS}")
+
+include_directories(
+ ${CUDA_INCLUDE_DIRS}
+)
+
+set(nvinfer_lib_name "nvinfer")
+set(nvinfer_plugin_lib_name "nvinfer_plugin")
+set(nvinfer_vc_plugin_lib_name "nvinfer_vc_plugin")
+set(nvonnxparser_lib_name "nvonnxparser")
+
+find_library_create_target(nvinfer ${nvinfer_lib_name} SHARED ${TRT_LIB_DIR})
+find_library(CUDART_LIB cudart_static HINTS ${CUDA_TOOLKIT_ROOT_DIR} PATH_SUFFIXES lib lib/x64 lib64)
+
+if (NOT MSVC)
+ find_library(RT_LIB rt)
+endif()
+
+set(CUDA_LIBRARIES ${CUDART_LIB})
+
+set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr -Xcompiler -Wno-deprecated-declarations")
+
+############################################################################################
+# TensorRT
+
+if(BUILD_SAMPLES)
+ add_subdirectory(samples)
+endif()

README.md

@@ -38,7 +38,7 @@ https://stackoverflow.com/questions/54429210/how-do-i-prevent-conda-from-activat
 ./build/sample_onnx_mnist -d ../data/tensorrt-sample-data/mnist/ --fp16
 # 5. 如何ssh遇到如下问题:
 ssh: connect to host github.com port 22: Connection refused
-Try this:
+Try this: https://stackoverflow.com/questions/7953806/github-ssh-via-public-wifi-port-22-blocked/45473512#45473512
 $ vim ~/.ssh/config and Add
 Host github.com
  Hostname ssh.github.com

cmake/modules/find_library_create_target.cmake

@@ -18,10 +18,13 @@
 macro(find_library_create_target target_name lib libtype hints)
  message(STATUS "========================= Importing and creating target ${target_name} ==========================")
  message(STATUS "Looking for library ${lib}")
+ message(STATUS "Library that was found ${${lib}_LIB_PATH}")
  if (CMAKE_BUILD_TYPE STREQUAL "Debug")
  find_library(${lib}_LIB_PATH ${lib}${TRT_DEBUG_POSTFIX} HINTS ${hints} NO_DEFAULT_PATH)
- endif()
+ message(STATUS "Library that was found ${${lib}_LIB_PATH}")
+ endif()
  find_library(${lib}_LIB_PATH ${lib} HINTS ${hints} NO_DEFAULT_PATH)
+ message(STATUS "Library that was found ${${lib}_LIB_PATH}")
  find_library(${lib}_LIB_PATH ${lib})
  message(STATUS "Library that was found ${${lib}_LIB_PATH}")
  add_library(${target_name} ${libtype} IMPORTED)

sampleMNIST.engine

@@ -92,7 +92,7 @@ inline bool parseArgs(Args& args, int32_t argc, char* argv[])
  {"useDLACore", required_argument, 0, 'u'}, {"batch", required_argument, 0, 'b'},
  {"timingCacheFile", required_argument, 0, 't'}, {nullptr, 0, nullptr, 0}};
  int32_t option_index = 0;
- arg = getopt_long(argc, argv, "hd:iu", long_options, &option_index);
+ arg = getopt_long(argc, argv, "hd:iu:o:s:", long_options, &option_index);
  if (arg == -1)
  {
  break;
@@ -103,7 +103,8 @@ inline bool parseArgs(Args& args, int32_t argc, char* argv[])
  case 'h': args.help = true; return true;
  case 'd':
  if (optarg)
- {
+ { 
+ std::cerr << "INFO: --datadir is valid" << std::endl;
  args.dataDirs.push_back(optarg);
  }
  else
@@ -116,12 +117,14 @@ inline bool parseArgs(Args& args, int32_t argc, char* argv[])
  if (optarg)
  {
  args.saveEngine = optarg;
+ std::cout << "parseArgs: save engine: " << args.saveEngine << std::endl;
  }
  break;
  case 'o':
  if (optarg)
  {
  args.loadEngine = optarg;
+ std::cerr << "INFO: --datadir is valid" << optarg << std::endl;
  }
  break;
  case 'i': args.runInInt8 = true; break;

samples/common/argsParser.h

@@ -16,6 +16,11 @@
 #
 SET(SAMPLE_SOURCES
  sampleOnnxMNIST.cpp
+ ../common/sampleEngines.cpp # original 
+ ../common/sampleDevice.cpp # below are save engine needed 
+ ../common/sampleOptions.cpp
+ ../common/sampleUtils.cpp
+ ../common/bfloat16.cpp
 )
 
 set(SAMPLE_PARSERS "onnx")

samples/sampleOnnxMNIST/CMakeLists.txt

@@ -33,6 +33,7 @@
 #include "common.h"
 #include "logger.h"
 #include "parserOnnxConfig.h"
+#include "sampleEngines.h"
 
 #include "NvInfer.h"
 #include <cuda_runtime_api.h>
@@ -46,14 +47,20 @@ using samplesCommon::SampleUniquePtr;
 
 const std::string gSampleName = "TensorRT.sample_onnx_mnist";
 
+struct sampleOnnxMNISTParams : samplesCommon::OnnxSampleParams
+{
+ std::string saveEngine;
+ std::string loadEngine;
+};
+
 //! \brief The SampleOnnxMNIST class implements the ONNX MNIST sample
 //!
 //! \details It creates the network using an ONNX model
 //!
 class SampleOnnxMNIST
 {
 public:
- SampleOnnxMNIST(const samplesCommon::OnnxSampleParams& params)
+ SampleOnnxMNIST(const sampleOnnxMNISTParams& params)
  : mParams(params)
  , mRuntime(nullptr)
  , mEngine(nullptr)
@@ -71,7 +78,7 @@ class SampleOnnxMNIST
  bool infer();
 
 private:
- samplesCommon::OnnxSampleParams mParams; //!< The parameters for the sample.
+ sampleOnnxMNISTParams mParams; //!< The parameters for the sample.
 
  nvinfer1::Dims mInputDims; //!< The dimensions of the input to the network.
  nvinfer1::Dims mOutputDims; //!< The dimensions of the output to the network.
@@ -108,6 +115,7 @@ class SampleOnnxMNIST
 //!
 bool SampleOnnxMNIST::build()
 {
+ // zqi tbd: builder的指针成员mImpl什么时候初始化的？//GPT: builder 的 mImpl 成员是在 createInferBuilder_INTERNAL 函数内部初始化的
  auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
  if (!builder)
  {
@@ -125,9 +133,9 @@ bool SampleOnnxMNIST::build()
  {
  return false;
  }
-
+ 
  auto parser
- = SampleUniquePtr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
+ = SampleUniquePtr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger())); // zqi: 关联 network和parser：https://github.com/onnx/onnx-tensorrt/blob/9c69a24bc2e20c8a511a4e6b06fd49639ec5300a/NvOnnxParser.cpp
  if (!parser)
  {
  return false;
@@ -149,7 +157,8 @@ bool SampleOnnxMNIST::build()
  }
  config->setProfileStream(*profileStream);
 
- SampleUniquePtr<IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};
+ // 生成了一个序列化的推理引擎数据，通常会被写入文件或传递给其他组件进行进一步处理或执行推理
+ SampleUniquePtr<IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)}; 
  if (!plan)
  {
  return false;
@@ -174,6 +183,13 @@ bool SampleOnnxMNIST::build()
  return false;
  }
 
+ // 使用 saveEngine 函数保存序列化的引擎
+ std::ofstream errStream; // 错误输出流（可以用 std::cerr 也可以用文件流）
+ std::string engineFileName = mParams.saveEngine; //"sampleMNIST.engine"; // 设置你想保存的文件名
+ if (!sample::saveEngine(*mEngine, engineFileName, std::cerr)) {
+ return false; // 保存失败，返回 false
+ }
+
  ASSERT(network->getNbInputs() == 1);
  mInputDims = network->getInput(0)->getDimensions();
  ASSERT(mInputDims.nbDims == 4);
@@ -353,9 +369,9 @@ bool SampleOnnxMNIST::verifyOutput(const samplesCommon::BufferManager& buffers)
 //!
 //! \brief Initializes members of the params struct using the command line args
 //!
-samplesCommon::OnnxSampleParams initializeSampleParams(const samplesCommon::Args& args)
+sampleOnnxMNISTParams initializeSampleParams(const samplesCommon::Args& args)
 {
- samplesCommon::OnnxSampleParams params;
+ sampleOnnxMNISTParams params;
  if (args.dataDirs.empty()) // Use default directories if user hasn't provided directory paths
  {
  params.dataDirs.push_back("data/mnist/");
@@ -373,6 +389,8 @@ samplesCommon::OnnxSampleParams initializeSampleParams(const samplesCommon::Args
  params.fp16 = args.runInFp16;
  params.bf16 = args.runInBf16;
  params.timingCacheFile = args.timingCacheFile;
+ params.saveEngine = args.saveEngine;
+ std::cout << "save Engine path: " << params.saveEngine << std::endl;
 
  return params;
 }