40 lines
1.2 KiB
C++
40 lines
1.2 KiB
C++
#ifndef ENTROPY_CALIBRATOR_H
|
|
#define ENTROPY_CALIBRATOR_H
|
|
|
|
#include <NvInfer.h>
|
|
#include <string>
|
|
#include <vector>
|
|
#include "macros.h"
|
|
|
|
//! \class Int8EntropyCalibrator2
|
|
//!
|
|
//! \brief Implements Entropy calibrator 2.
|
|
//! CalibrationAlgoType is kENTROPY_CALIBRATION_2.
|
|
//!
|
|
class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 {
|
|
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 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_;
|
|
int input_w_;
|
|
int input_h_;
|
|
int img_idx_;
|
|
std::string img_dir_;
|
|
std::vector<std::string> img_files_;
|
|
size_t input_count_;
|
|
std::string calib_table_name_;
|
|
const char* input_blob_name_;
|
|
bool read_cache_;
|
|
void* device_input_;
|
|
std::vector<char> calib_cache_;
|
|
};
|
|
|
|
#endif // ENTROPY_CALIBRATOR_H
|