initial commit

wangrong

2025-01-03 2c4b8838f89bebca97615c5f0bac2c58e0349771

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
#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