From b7a754e2de7b21989a868f2740b82d8417dad1cf Mon Sep 17 00:00:00 2001 From: Aymeric Dujardin Date: Mon, 5 Jul 2021 06:02:56 +0200 Subject: [PATCH] Fix for TensorRT 8 support (#540) --- yolov5/calibrator.cpp | 8 +++--- yolov5/calibrator.h | 14 +++++++--- yolov5/logging.h | 8 +++++- yolov5/yololayer.cu | 44 ++++++++++++++--------------- yolov5/yololayer.h | 65 ++++++++++++++++++++++++------------------- 5 files changed, 80 insertions(+), 59 deletions(-) diff --git a/yolov5/calibrator.cpp b/yolov5/calibrator.cpp index 1edcde1..2e2f586 100644 --- a/yolov5/calibrator.cpp +++ b/yolov5/calibrator.cpp @@ -26,12 +26,12 @@ 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_ > (int)img_files_.size()) { return false; @@ -57,7 +57,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(); @@ -71,7 +71,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); diff --git a/yolov5/calibrator.h b/yolov5/calibrator.h index 1cc9dbb..27f2d37 100644 --- a/yolov5/calibrator.h +++ b/yolov5/calibrator.h @@ -5,6 +5,12 @@ #include #include +#if NV_TENSORRT_MAJOR >= 8 +#define TRT_NOEXCEPT noexcept +#else +#define TRT_NOEXCEPT +#endif + //! \class Int8EntropyCalibrator2 //! //! \brief Implements Entropy calibrator 2. @@ -16,10 +22,10 @@ public: Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, 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_; diff --git a/yolov5/logging.h b/yolov5/logging.h index 602b69f..1339ee2 100644 --- a/yolov5/logging.h +++ b/yolov5/logging.h @@ -26,6 +26,12 @@ #include #include +#if NV_TENSORRT_MAJOR >= 8 +#define TRT_NOEXCEPT noexcept +#else +#define TRT_NOEXCEPT +#endif + using Severity = nvinfer1::ILogger::Severity; class LogStreamConsumerBuffer : public std::stringbuf @@ -236,7 +242,7 @@ public: //! 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/yolov5/yololayer.cu b/yolov5/yololayer.cu index 525bf8d..4c042a6 100644 --- a/yolov5/yololayer.cu +++ b/yolov5/yololayer.cu @@ -78,7 +78,7 @@ namespace nvinfer1 assert(d == a + length); } - void YoloLayerPlugin::serialize(void* buffer) const + void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT { using namespace Tn; char* d = static_cast(buffer), *a = d; @@ -95,17 +95,17 @@ namespace nvinfer1 assert(d == a + getSerializationSize()); } - size_t YoloLayerPlugin::getSerializationSize() const + size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT { return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV5NetWidth) + sizeof(mYoloV5NetHeight) + sizeof(mMaxOutObject); } - int YoloLayerPlugin::initialize() + int YoloLayerPlugin::initialize() TRT_NOEXCEPT { return 0; } - Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) + Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT { //output the result to channel int totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float); @@ -114,63 +114,63 @@ namespace nvinfer1 } // Set plugin namespace - void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) + void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT { mPluginNamespace = pluginNamespace; } - const char* YoloLayerPlugin::getPluginNamespace() const + const char* YoloLayerPlugin::getPluginNamespace() const TRT_NOEXCEPT { return mPluginNamespace; } // Return the DataType of the plugin output at the requested index - DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const + DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT { return DataType::kFLOAT; } // Return true if output tensor is broadcast across a batch. - bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const + bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT { return false; } // Return true if plugin can use input that is broadcast across batch without replication. - bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const + bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT { return false; } - void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) + void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT { } // Attach the plugin object to an execution context and grant the plugin the access to some context resource. - void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) + void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT { } // Detach the plugin object from its execution context. - void YoloLayerPlugin::detachFromContext() {} + void YoloLayerPlugin::detachFromContext() TRT_NOEXCEPT {} - const char* YoloLayerPlugin::getPluginType() const + const char* YoloLayerPlugin::getPluginType() const TRT_NOEXCEPT { return "YoloLayer_TRT"; } - const char* YoloLayerPlugin::getPluginVersion() const + const char* YoloLayerPlugin::getPluginVersion() const TRT_NOEXCEPT { return "1"; } - void YoloLayerPlugin::destroy() + void YoloLayerPlugin::destroy() TRT_NOEXCEPT { delete this; } // Clone the plugin - IPluginV2IOExt* YoloLayerPlugin::clone() const + IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT { YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, mYoloKernel); p->setPluginNamespace(mPluginNamespace); @@ -252,7 +252,7 @@ namespace nvinfer1 } - int YoloLayerPlugin::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) + int YoloLayerPlugin::enqueue(int batchSize, const void* const* inputs, void* TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT { forwardGpu((const float* const*)inputs, (float*)outputs[0], stream, batchSize); return 0; @@ -269,22 +269,22 @@ namespace nvinfer1 mFC.fields = mPluginAttributes.data(); } - const char* YoloPluginCreator::getPluginName() const + const char* YoloPluginCreator::getPluginName() const TRT_NOEXCEPT { return "YoloLayer_TRT"; } - const char* YoloPluginCreator::getPluginVersion() const + const char* YoloPluginCreator::getPluginVersion() const TRT_NOEXCEPT { return "1"; } - const PluginFieldCollection* YoloPluginCreator::getFieldNames() + const PluginFieldCollection* YoloPluginCreator::getFieldNames() TRT_NOEXCEPT { return &mFC; } - IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) + IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT { assert(fc->nbFields == 2); assert(strcmp(fc->fields[0].name, "netinfo") == 0); @@ -301,7 +301,7 @@ namespace nvinfer1 return obj; } - IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) + IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT { // This object will be deleted when the network is destroyed, which will // call YoloLayerPlugin::destroy() diff --git a/yolov5/yololayer.h b/yolov5/yololayer.h index 49f6474..7224bde 100644 --- a/yolov5/yololayer.h +++ b/yolov5/yololayer.h @@ -5,6 +5,15 @@ #include #include "NvInfer.h" + +#if NV_TENSORRT_MAJOR >= 8 +#define TRT_NOEXCEPT noexcept +#define TRT_CONST_ENQUEUE const +#else +#define TRT_NOEXCEPT +#define TRT_CONST_ENQUEUE +#endif + namespace Yolo { static constexpr int CHECK_COUNT = 3; @@ -38,53 +47,53 @@ namespace nvinfer1 YoloLayerPlugin(const void* data, size_t length); ~YoloLayerPlugin(); - int getNbOutputs() const override + int getNbOutputs() const TRT_NOEXCEPT override { return 1; } - Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override; + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT override; - int initialize() override; + int initialize() TRT_NOEXCEPT override; - virtual void terminate() override {}; + virtual void terminate() TRT_NOEXCEPT override {}; - virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0; } + virtual size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { return 0; } - virtual int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override; + virtual int enqueue(int batchSize, const void* const* inputs, void*TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT override; - virtual size_t getSerializationSize() const override; + virtual size_t getSerializationSize() const TRT_NOEXCEPT override; - virtual void serialize(void* buffer) const override; + virtual void serialize(void* buffer) const TRT_NOEXCEPT override; - bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override { + bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const TRT_NOEXCEPT override { return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; } - const char* getPluginType() const override; + const char* getPluginType() const TRT_NOEXCEPT override; - const char* getPluginVersion() const override; + const char* getPluginVersion() const TRT_NOEXCEPT override; - void destroy() override; + void destroy() TRT_NOEXCEPT override; - IPluginV2IOExt* clone() const override; + IPluginV2IOExt* clone() const TRT_NOEXCEPT override; - void setPluginNamespace(const char* pluginNamespace) override; + void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override; - const char* getPluginNamespace() const override; + const char* getPluginNamespace() const TRT_NOEXCEPT override; - DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override; + DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override; - bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override; + bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT override; - bool canBroadcastInputAcrossBatch(int inputIndex) const override; + bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override; void attachToContext( - cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override; + cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT override; - void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override; + void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT override; - void detachFromContext() override; + void detachFromContext() TRT_NOEXCEPT override; private: void forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize = 1); @@ -106,22 +115,22 @@ namespace nvinfer1 ~YoloPluginCreator() override = default; - const char* getPluginName() const override; + const char* getPluginName() const TRT_NOEXCEPT override; - const char* getPluginVersion() const override; + const char* getPluginVersion() const TRT_NOEXCEPT override; - const PluginFieldCollection* getFieldNames() override; + const PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override; - IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override; + IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT override; - IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override; + IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT override; - void setPluginNamespace(const char* libNamespace) override + void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override { mNamespace = libNamespace; } - const char* getPluginNamespace() const override + const char* getPluginNamespace() const TRT_NOEXCEPT override { return mNamespace.c_str(); }