detr添加对trt8的支持 (#1457)

* detr添加对trt8的支持

* 修改格式

* 修改格式

* 修改格式

* 修改格式

* Update detr.cpp

---------

Co-authored-by: Wang Xinyu <shaywxy@gmail.com>

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")
 

detr/calibrator.hpp

@@ -8,7 +8,8 @@
 #include <fstream>
 #include <algorithm>
 #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<int>(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<const char*>(cache), length);

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));

detr/logging.h

@@ -26,6 +26,9 @@
 #include <sstream>
 #include <string>
 
+#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;
     }
 

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