From e4da5a8367c685456a2aa3a9d30f6f502b1eae25 Mon Sep 17 00:00:00 2001 From: Rex <74702576+Rex-LK@users.noreply.github.com> Date: Mon, 18 Mar 2024 11:38:38 +0800 Subject: [PATCH] =?UTF-8?q?detr=E6=B7=BB=E5=8A=A0=E5=AF=B9trt8=E7=9A=84?= =?UTF-8?q?=E6=94=AF=E6=8C=81=20(#1457)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * detr添加对trt8的支持 * 修改格式 * 修改格式 * 修改格式 * 修改格式 * Update detr.cpp --------- Co-authored-by: Wang Xinyu --- detr/CMakeLists.txt | 8 ++++---- detr/calibrator.hpp | 19 ++++++++++--------- detr/detr.cpp | 10 ++++++++-- detr/logging.h | 5 ++++- detr/macros.h | 29 +++++++++++++++++++++++++++++ 5 files changed, 55 insertions(+), 16 deletions(-) create mode 100644 detr/macros.h diff --git a/detr/CMakeLists.txt b/detr/CMakeLists.txt index 2b4e145..2b23ece 100644 --- a/detr/CMakeLists.txt +++ b/detr/CMakeLists.txt @@ -13,11 +13,11 @@ find_package(CUDA REQUIRED) include_directories(${PROJECT_SOURCE_DIR}/include) # include and link dirs of cuda and tensorrt, you need adapt them if yours are different # cuda -include_directories(/usr/local/cuda-10.2/include) -link_directories(/usr/local/cuda-10.2/lib64) +include_directories(/usr/local/cuda/include) +link_directories(/usr/local/cuda/lib64) # tensorrt -include_directories(/home/jushi/TensorRT-7.2.1.6/include) -link_directories(/home/jushi/TensorRT-7.2.1.6/lib) +include_directories(/data/app/TensorRT-8.4.3.1/include) +link_directories(/data/app/TensorRT-8.4.3.1/lib) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") diff --git a/detr/calibrator.hpp b/detr/calibrator.hpp index 4d20f06..9fff549 100644 --- a/detr/calibrator.hpp +++ b/detr/calibrator.hpp @@ -8,7 +8,8 @@ #include #include #include "common.hpp" - +#include "macros.h" + //! \class Int8EntropyCalibrator2 //! //! \brief Implements Entropy calibrator 2. @@ -21,10 +22,10 @@ class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 { const char* input_blob_name, bool read_cache = true); virtual ~Int8EntropyCalibrator2(); - int getBatchSize() const override; - bool getBatch(void* bindings[], const char* names[], int nbBindings) override; - const void* readCalibrationCache(size_t& length) override; - void writeCalibrationCache(const void* cache, size_t length) override; + int getBatchSize() const TRT_NOEXCEPT override; + bool getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT override; + const void* readCalibrationCache(size_t& length) TRT_NOEXCEPT override; + void writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT override; private: int batchsize_; @@ -62,11 +63,11 @@ Int8EntropyCalibrator2::~Int8EntropyCalibrator2() { CUDA_CHECK(cudaFree(device_input_)); } -int Int8EntropyCalibrator2::getBatchSize() const { +int Int8EntropyCalibrator2::getBatchSize() const TRT_NOEXCEPT { return batchsize_; } -bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings) { +bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT { if (img_idx_ + batchsize_ > static_cast(img_files_.size())) { return false; } @@ -97,7 +98,7 @@ bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int return true; } -const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) { +const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) TRT_NOEXCEPT { std::cout << "reading calib cache: " << calib_table_name_ << std::endl; calib_cache_.clear(); std::ifstream input(calib_table_name_, std::ios::binary); @@ -109,7 +110,7 @@ const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) { return length ? calib_cache_.data() : nullptr; } -void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length) { +void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT { std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl; std::ofstream output(calib_table_name_, std::ios::binary); output.write(reinterpret_cast(cache), length); diff --git a/detr/detr.cpp b/detr/detr.cpp index a7bc476..443c8c6 100644 --- a/detr/detr.cpp +++ b/detr/detr.cpp @@ -170,8 +170,11 @@ int num_heads = 8 v_shuffle->setName((lname + ".v_shuffle").c_str()); v_shuffle->setReshapeDimensions(Dims3{ -1, num_heads, head_dim }); v_shuffle->setSecondTranspose(Permutation{ 1, 0, 2 }); - +#if NV_TENSORRT_MAJOR >= 8 + auto q_product_k = network->addMatrixMultiply(*q_shuffle->getOutput(0), nvinfer1::MatrixOperation::kNONE, *k_shuffle->getOutput(0), nvinfer1::MatrixOperation::kTRANSPOSE); +#else auto q_product_k = network->addMatrixMultiply(*q_shuffle->getOutput(0), false, *k_shuffle->getOutput(0), true); +#endif assert(q_product_k); // src_key_padding_mask are all false, so do nothing here @@ -180,8 +183,11 @@ int num_heads = 8 auto softmax = network->addSoftMax(*q_product_k->getOutput(0)); assert(softmax); softmax->setAxes(4); - +#if NV_TENSORRT_MAJOR >= 8 + auto attn_product_v = network->addMatrixMultiply(*softmax->getOutput(0), nvinfer1::MatrixOperation::kNONE, *v_shuffle->getOutput(0), nvinfer1::MatrixOperation::kNONE); +#else auto attn_product_v = network->addMatrixMultiply(*softmax->getOutput(0), false, *v_shuffle->getOutput(0), false); +#endif assert(attn_product_v); auto attn_shuffle = network->addShuffle(*attn_product_v->getOutput(0)); diff --git a/detr/logging.h b/detr/logging.h index c9c9d6c..4f300b1 100644 --- a/detr/logging.h +++ b/detr/logging.h @@ -26,6 +26,9 @@ #include #include +#include "macros.h" + + using Severity = nvinfer1::ILogger::Severity; class LogStreamConsumerBuffer : public std::stringbuf { @@ -208,7 +211,7 @@ class Logger : public nvinfer1::ILogger { //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the //! inheritance from nvinfer1::ILogger //! - void log(Severity severity, const char* msg) override { + void log(Severity severity, const char* msg) TRT_NOEXCEPT override { LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; } diff --git a/detr/macros.h b/detr/macros.h new file mode 100644 index 0000000..92d9edd --- /dev/null +++ b/detr/macros.h @@ -0,0 +1,29 @@ +#ifndef __MACROS_H +#define __MACROS_H + +#include "NvInfer.h" + +#ifdef API_EXPORTS +#if defined(_MSC_VER) +#define API __declspec(dllexport) +#else +#define API __attribute__((visibility("default"))) +#endif +#else + +#if defined(_MSC_VER) +#define API __declspec(dllimport)a +#else +#define API +#endif +#endif // API_EXPORTS + +#if NV_TENSORRT_MAJOR >= 8 +#define TRT_NOEXCEPT noexcept +#define TRT_CONST_ENQUEUE const +#else +#define TRT_NOEXCEPT +#define TRT_CONST_ENQUEUE +#endif + +#endif // __MACROS_H