From b6a7ec5d15266b8da47a4c34ce6a786609c3cbd9 Mon Sep 17 00:00:00 2001 From: unknown Date: Thu, 25 Nov 2021 23:18:31 -0300 Subject: [PATCH] Fix INT8 calibrator --- nvdsinfer_custom_impl_Yolo/calibrator.cpp | 10 +++--- nvdsinfer_custom_impl_Yolo/calibrator.h | 14 ++++---- nvdsinfer_custom_impl_Yolo/yolo.cpp | 39 ++++++++++++----------- 3 files changed, 32 insertions(+), 31 deletions(-) diff --git a/nvdsinfer_custom_impl_Yolo/calibrator.cpp b/nvdsinfer_custom_impl_Yolo/calibrator.cpp index b335cb8..f5a4408 100644 --- a/nvdsinfer_custom_impl_Yolo/calibrator.cpp +++ b/nvdsinfer_custom_impl_Yolo/calibrator.cpp @@ -9,7 +9,7 @@ namespace nvinfer1 { - int8EntroyCalibrator::int8EntroyCalibrator(const int &batchsize, const int &channels, const int &height, const int &width, const int &letterbox, const std::string &imgPath, + Int8EntropyCalibrator2::Int8EntropyCalibrator2(const int &batchsize, const int &channels, const int &height, const int &width, const int &letterbox, const std::string &imgPath, const std::string &calibTablePath):batchSize(batchsize), inputC(channels), inputH(height), inputW(width), letterBox(letterbox), calibTablePath(calibTablePath), imageIndex(0) { inputCount = batchsize * channels * height * width; @@ -23,14 +23,14 @@ namespace nvinfer1 CUDA_CHECK(cudaMalloc(&deviceInput, inputCount * sizeof(float))); } - int8EntroyCalibrator::~int8EntroyCalibrator() + Int8EntropyCalibrator2::~Int8EntropyCalibrator2() { CUDA_CHECK(cudaFree(deviceInput)); if (batchData) delete[] batchData; } - bool int8EntroyCalibrator::getBatch(void **bindings, const char **names, int nbBindings) + bool Int8EntropyCalibrator2::getBatch(void **bindings, const char **names, int nbBindings) noexcept { if (imageIndex + batchSize > uint(imgPaths.size())) return false; @@ -54,7 +54,7 @@ namespace nvinfer1 return true; } - const void* int8EntroyCalibrator::readCalibrationCache(std::size_t &length) + const void* Int8EntropyCalibrator2::readCalibrationCache(std::size_t &length) noexcept { calibrationCache.clear(); std::ifstream input(calibTablePath, std::ios::binary); @@ -68,7 +68,7 @@ namespace nvinfer1 return length ? calibrationCache.data() : nullptr; } - void int8EntroyCalibrator::writeCalibrationCache(const void *cache, std::size_t length) + void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, std::size_t length) noexcept { std::ofstream output(calibTablePath, std::ios::binary); output.write(reinterpret_cast(cache), length); diff --git a/nvdsinfer_custom_impl_Yolo/calibrator.h b/nvdsinfer_custom_impl_Yolo/calibrator.h index a78e062..f912761 100644 --- a/nvdsinfer_custom_impl_Yolo/calibrator.h +++ b/nvdsinfer_custom_impl_Yolo/calibrator.h @@ -24,9 +24,9 @@ #endif namespace nvinfer1 { - class int8EntroyCalibrator : public nvinfer1::IInt8EntropyCalibrator2 { + class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 { public: - int8EntroyCalibrator(const int &batchsize, + Int8EntropyCalibrator2(const int &batchsize, const int &channels, const int &height, const int &width, @@ -34,11 +34,11 @@ namespace nvinfer1 { const std::string &imgPath, const std::string &calibTablePath); - virtual ~int8EntroyCalibrator(); - int getBatchSize() const override { return batchSize; } - bool getBatch(void *bindings[], const char *names[], int nbBindings) override; - const void *readCalibrationCache(std::size_t &length) override; - void writeCalibrationCache(const void *ptr, std::size_t length) override; + virtual ~Int8EntropyCalibrator2(); + int getBatchSize() const noexcept override; + bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override; + const void* readCalibrationCache(std::size_t& length) noexcept override; + void writeCalibrationCache(const void* cache, size_t length) noexcept override; private: int batchSize; diff --git a/nvdsinfer_custom_impl_Yolo/yolo.cpp b/nvdsinfer_custom_impl_Yolo/yolo.cpp index b96d7ed..8e7d9e0 100644 --- a/nvdsinfer_custom_impl_Yolo/yolo.cpp +++ b/nvdsinfer_custom_impl_Yolo/yolo.cpp @@ -73,6 +73,23 @@ nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder) parseConfigBlocks(); orderParams(&m_OutputMasks); + std::vector weights = loadWeights(m_WtsFilePath, m_NetworkType); + std::vector trtWeights; + + nvinfer1::INetworkDefinition *network = builder->createNetworkV2(0); + if (parseModel(*network) != NVDSINFER_SUCCESS) { + network->destroy(); + return nullptr; + } + + std::cout << "Building the TensorRT Engine" << std::endl; + + if (m_LetterBox == 1) { + std::cout << "\nNOTE: letter_box is set in cfg file, make sure to set maintain-aspect-ratio=1 in config_infer file to get better accuracy\n" << std::endl; + } + + nvinfer1::IBuilderConfig *config = builder->createBuilderConfig(); + if (m_NetworkMode == "INT8" && !fileExists(m_Int8CalibPath)) { assert(builder->platformHasFastInt8()); #ifdef OPENCV @@ -92,31 +109,15 @@ nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder) std::cerr << "INT8_CALIB_BATCH_SIZE not set" << std::endl; std::abort(); } - nvinfer1::int8EntroyCalibrator *calibrator = new nvinfer1::int8EntroyCalibrator(calib_batch_size, m_InputC, m_InputH, m_InputW, m_LetterBox, calib_image_list, m_Int8CalibPath); - builder->setInt8Mode(true); - builder->setInt8Calibrator(calibrator); + nvinfer1::Int8EntropyCalibrator2 *calibrator = new nvinfer1::Int8EntropyCalibrator2(calib_batch_size, m_InputC, m_InputH, m_InputW, m_LetterBox, calib_image_list, m_Int8CalibPath); + config->setFlag(nvinfer1::BuilderFlag::kINT8); + config->setInt8Calibrator(calibrator); #else std::cerr << "OpenCV is required to run INT8 calibrator" << std::endl; std::abort(); #endif } - std::vector weights = loadWeights(m_WtsFilePath, m_NetworkType); - std::vector trtWeights; - - nvinfer1::INetworkDefinition *network = builder->createNetworkV2(0); - if (parseModel(*network) != NVDSINFER_SUCCESS) { - network->destroy(); - return nullptr; - } - - std::cout << "Building the TensorRT Engine" << std::endl; - - if (m_LetterBox == 1) { - std::cout << "\nNOTE: letter_box is set in cfg file, make sure to set maintain-aspect-ratio=1 in config_infer file to get better accuracy\n" << std::endl; - } - - nvinfer1::IBuilderConfig *config = builder->createBuilderConfig(); nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config); if (engine) { std::cout << "Building complete\n" << std::endl;