mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

Fix the INT8 calibrator sintax

Browse Source 
Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>
This commit is contained in:
Davide Sapienza
2020-04-08 10:22:03 +02:00
parent c36befaf2b
commit 5562f599a6
4 changed files with 42 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/Int8Calibrator.h
+11 -3
View File
@@ -18,9 +18,17 @@
#include "tkdnn.h"
#include "utils.h"
class Int8EntropyCalibrator : public nvinfer1::IInt8EntropyCalibrator{
/*
* Int8EntropyCalibrator implements the INT8 calibrator to achieve the
* INT8 quantization. It uses a BatchStream stream to scroll through
* images data. It also implements the calibration cache, a way to
* save the calibration process results to reduce the running time:
* the calibration process takes a long time.
*/
class Int8EntropyCalibrator : public nvinfer1::IInt8EntropyCalibrator {
public:
Int8EntropyCalibrator(BatchStream& stream, int firstBatch, const std::string& calibTableFilePath, const std::string& inputBlobName, bool readCache = true);
Int8EntropyCalibrator(BatchStream& stream, int firstBatch, const std::string& calibTableFilePath,
const std::string& inputBlobName, bool readCache = true);
virtual ~Int8EntropyCalibrator() { checkCuda(cudaFree(mDeviceInput)); }
int getBatchSize() const override { return mStream.getBatchSize(); }
bool getBatch(void* bindings[], const char* names[], int nbBindings) override;
@@ -29,7 +37,7 @@ public:
private:
BatchStream mStream;
const std::string mCalibTableFilePath{nullptr};
const std::string mCalibTableFilePath{ nullptr };
const std::string mInputBlobName;
bool mReadCache{ true };