34 lines
1.5 KiB
C++
34 lines
1.5 KiB
C++
#pragma once
|
|
#include "NvInfer.h"
|
|
#include "engine_api.h" // ENGINE_API (__declspec dllexport / dllimport)
|
|
|
|
// Class used for int8 calibration
|
|
class ENGINE_API Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 {
|
|
public:
|
|
Int8EntropyCalibrator2(int32_t batchSize, int32_t inputW, int32_t inputH, const std::string &calibDataDirPath,
|
|
const std::string &calibTableName, const std::string &inputBlobName,
|
|
const std::array<float, 3> &subVals = {0.f, 0.f, 0.f}, const std::array<float, 3> &divVals = {1.f, 1.f, 1.f},
|
|
bool normalize = true, bool readCache = true);
|
|
virtual ~Int8EntropyCalibrator2();
|
|
// Abstract base class methods which must be implemented
|
|
int32_t getBatchSize() const noexcept override;
|
|
bool getBatch(void *bindings[], char const *names[], int32_t nbBindings) noexcept override;
|
|
void const *readCalibrationCache(std::size_t &length) noexcept override;
|
|
void writeCalibrationCache(void const *ptr, std::size_t length) noexcept override;
|
|
|
|
private:
|
|
const int32_t m_batchSize;
|
|
const int32_t m_inputW;
|
|
const int32_t m_inputH;
|
|
int32_t m_imgIdx;
|
|
std::vector<std::string> m_imgPaths;
|
|
size_t m_inputCount;
|
|
const std::string m_calibTableName;
|
|
const std::string m_inputBlobName;
|
|
const std::array<float, 3> m_subVals;
|
|
const std::array<float, 3> m_divVals;
|
|
const bool m_normalize;
|
|
const bool m_readCache;
|
|
void *m_deviceInput;
|
|
std::vector<char> m_calibCache;
|
|
}; |