thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix INT8 calibrator

Browse Source
This commit is contained in:
unknown
2021-11-25 23:18:31 -03:00
parent af8009d640
commit b6a7ec5d15
3 changed files with 32 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
nvdsinfer_custom_impl_Yolo/calibrator.cpp
View File

@@ -9,7 +9,7 @@
namespace nvinfer1 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) const std::string &calibTablePath):batchSize(batchsize), inputC(channels), inputH(height), inputW(width), letterBox(letterbox), calibTablePath(calibTablePath), imageIndex(0)
{ {
inputCount = batchsize * channels * height * width; inputCount = batchsize * channels * height * width;
@@ -23,14 +23,14 @@ namespace nvinfer1
CUDA_CHECK(cudaMalloc(&deviceInput, inputCount * sizeof(float))); CUDA_CHECK(cudaMalloc(&deviceInput, inputCount * sizeof(float)));
} }
int8EntroyCalibrator::~int8EntroyCalibrator() Int8EntropyCalibrator2::~Int8EntropyCalibrator2()
{ {
CUDA_CHECK(cudaFree(deviceInput)); CUDA_CHECK(cudaFree(deviceInput));
if (batchData) if (batchData)
delete[] 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())) if (imageIndex + batchSize > uint(imgPaths.size()))
return false; return false;
@@ -54,7 +54,7 @@ namespace nvinfer1
return true; return true;
} }
const void* int8EntroyCalibrator::readCalibrationCache(std::size_t &length) const void* Int8EntropyCalibrator2::readCalibrationCache(std::size_t &length) noexcept
{ {
calibrationCache.clear(); calibrationCache.clear();
std::ifstream input(calibTablePath, std::ios::binary); std::ifstream input(calibTablePath, std::ios::binary);
@@ -68,7 +68,7 @@ namespace nvinfer1
return length ? calibrationCache.data() : nullptr; 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); std::ofstream output(calibTablePath, std::ios::binary);
output.write(reinterpret_cast<const char*>(cache), length); output.write(reinterpret_cast<const char*>(cache), length);

14
nvdsinfer_custom_impl_Yolo/calibrator.h
View File

@@ -24,9 +24,9 @@
#endif #endif
namespace nvinfer1 { namespace nvinfer1 {
class int8EntroyCalibrator : public nvinfer1::IInt8EntropyCalibrator2 { class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 {
public: public:
int8EntroyCalibrator(const int &batchsize, Int8EntropyCalibrator2(const int &batchsize,
const int &channels, const int &channels,
const int &height, const int &height,
const int &width, const int &width,
@@ -34,11 +34,11 @@ namespace nvinfer1 {
const std::string &imgPath, const std::string &imgPath,
const std::string &calibTablePath); const std::string &calibTablePath);
virtual ~int8EntroyCalibrator(); virtual ~Int8EntropyCalibrator2();
int getBatchSize() const override { return batchSize; } int getBatchSize() const noexcept override;
bool getBatch(void *bindings[], const char *names[], int nbBindings) override; bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override;
const void *readCalibrationCache(std::size_t &length) override; const void* readCalibrationCache(std::size_t& length) noexcept override;
void writeCalibrationCache(const void *ptr, std::size_t length) override; void writeCalibrationCache(const void* cache, size_t length) noexcept override;
private: private:
int batchSize; int batchSize;

39
nvdsinfer_custom_impl_Yolo/yolo.cpp
View File

@@ -73,6 +73,23 @@ nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder)
parseConfigBlocks(); parseConfigBlocks();
orderParams(&m_OutputMasks); orderParams(&m_OutputMasks);
std::vector<float> weights = loadWeights(m_WtsFilePath, m_NetworkType);
std::vector<nvinfer1::Weights> 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)) { if (m_NetworkMode == "INT8" && !fileExists(m_Int8CalibPath)) {
assert(builder->platformHasFastInt8()); assert(builder->platformHasFastInt8());
#ifdef OPENCV #ifdef OPENCV
@@ -92,31 +109,15 @@ nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder)
std::cerr << "INT8_CALIB_BATCH_SIZE not set" << std::endl; std::cerr << "INT8_CALIB_BATCH_SIZE not set" << std::endl;
std::abort(); std::abort();
} }
nvinfer1::int8EntroyCalibrator *calibrator = new nvinfer1::int8EntroyCalibrator(calib_batch_size, m_InputC, m_InputH, m_InputW, m_LetterBox, calib_image_list, m_Int8CalibPath); nvinfer1::Int8EntropyCalibrator2 *calibrator = new nvinfer1::Int8EntropyCalibrator2(calib_batch_size, m_InputC, m_InputH, m_InputW, m_LetterBox, calib_image_list, m_Int8CalibPath);
builder->setInt8Mode(true); config->setFlag(nvinfer1::BuilderFlag::kINT8);
builder->setInt8Calibrator(calibrator); config->setInt8Calibrator(calibrator);
#else #else
std::cerr << "OpenCV is required to run INT8 calibrator" << std::endl; std::cerr << "OpenCV is required to run INT8 calibrator" << std::endl;
std::abort(); std::abort();
#endif #endif
} }
std::vector<float> weights = loadWeights(m_WtsFilePath, m_NetworkType);
std::vector<nvinfer1::Weights> 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); nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
if (engine) { if (engine) {
std::cout << "Building complete\n" << std::endl; std::cout << "Building complete\n" << std::endl;