experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Use IInt8EntropyCalibrator2 for TRT 5.1 onward

Browse Source
This commit is contained in:
Trevor Morris 2019-02-14 16:33:56 -08:00
parent b51e216294
commit 3c4a320722
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
View File

@ -34,7 +34,12 @@ namespace tensorrt {
// TRTs pull model for calibration. When TRT implements a means for
// a push calibration This class should be updated accordingly
// IInt8EntropyCalibrator2 is prefferred for TRT 5.1+.
#if NV_TENSORRT_MAJOR > 5 || (NV_TENSORRT_MAJOR == 5 && NV_TENSORRT_MINOR >= 1)
struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator2 {
#else
struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator {
#endif
public:
// Construct a calibrator for future calibration.
TRTInt8Calibrator(