Add the Int8 calibrator and the tensorRT Int8 inference

Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>

include/tkDNN/Int8BatchStream.h

@@ -0,0 +1,65 @@
+#ifndef INT8BATCHSTREAM_H
+#define INT8BATCHSTREAM_H
+
+#include <vector>
+#include <assert.h>
+#include <algorithm>
+#include <iterator>
+#include <stdint.h>
+#include <iostream>
+#include <string>
+#include "NvInfer.h"
+
+#include <fstream>
+#include <iomanip>
+
+#include <opencv2/core/core.hpp>
+#include <opencv2/dnn/dnn.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+
+#include "tkdnn.h"
+#include "utils.h"
+
+class BatchStream
+{
+public:
+	BatchStream(tk::dnn::dataDim_t dim, int batchSize, int maxBatches, const std::string& fileimglist, const std::string& filelabellist);
+	virtual ~BatchStream() {}
+	void reset(int firstBatch);
+	bool next();
+	void skip(int skipCount);
+	float *getBatch() { return mBatch.data();}
+	float *getLabels() { return mLabels.data();}
+	int getBatchesRead() const { return mBatchCount; }
+	int getBatchSize() const { return mBatchSize; }
+	nvinfer1::DimsNCHW getDims() const { return mDims; }
+	float* getFileBatch() { return &mFileBatch[0]; }
+	float* getFileLabels() { return &mFileLabels[0]; }
+	void readInListFile(const std::string& dataFilePath, std::vector<std::string>& mListIn);
+	void readCVimage(std::string inputFileName, std::vector<float>& res, bool fixshape = true);
+	void readLabels(std::string inputFileName ,std::vector<float>& ris);
+	bool update();
+
+private:
+	int mBatchSize{ 0 };
+	int mMaxBatches{ 0 };
+	int mBatchCount{ 0 };
+	int mFileCount{ 0 }, mFileBatchPos{ 0 };
+	int mImageSize{ 0 };
+
+	nvinfer1::DimsNCHW mDims;
+	std::vector<float> mBatch;
+	std::vector<float> mLabels;
+	std::vector<float> mFileBatch;
+	std::vector<float> mFileLabels;
+
+	int mHeight;
+	int mWidth;
+	std::string mFileImgList;
+	std::vector<std::string> mListImg;
+	std::string mFileLabelList;
+	std::vector<std::string> mListLabel;
+};
+
+#endif //INT8BATCHSTREAM

include/tkDNN/Int8Calibrator.h

@@ -0,0 +1,41 @@
+#ifndef INT8CALIBRATOR_H
+#define INT8CALIBRATOR_H
+
+#include <vector>
+#include <assert.h>
+#include <algorithm>
+#include <iterator>
+#include <stdint.h>
+#include <iostream>
+#include <string>
+#include "NvInfer.h"
+
+#include <fstream>
+#include <iomanip>
+
+#include "Int8BatchStream.h"
+
+#include "tkdnn.h"
+#include "utils.h"
+
+class Int8EntropyCalibrator : public nvinfer1::IInt8EntropyCalibrator{
+public:
+	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;
+	const void* readCalibrationCache(size_t& length) override;
+	void writeCalibrationCache(const void* cache, size_t length) override;
+
+private:
+	BatchStream mStream;
+	const std::string mCalibTableFilePath{nullptr};
+	const std::string mInputBlobName;
+	bool mReadCache{ true };
+
+	size_t mInputCount;
+	void* mDeviceInput{ nullptr };
+	std::vector<char> mCalibrationCache;
+};
+
+#endif //INT8CALIBRATOR_H

include/tkDNN/NetworkRT.h

@@ -32,7 +32,6 @@ using namespace nvinfer1;
 #include "pluginsRT/YoloRT.h"
 #include "pluginsRT/UpsampleRT.h"
 #include "pluginsRT/ResizeLayerRT.h"
-//#include "pluginsRT/Int8Calibrator.h"
 #include "pluginsRT/DeformableConvRT.h"
 #include "pluginsRT/FlattenConcatRT.h"
 #include "pluginsRT/ReshapeRT.h"
@@ -56,6 +55,7 @@ public:
     nvinfer1::IBuilder *builderRT;
     nvinfer1::IRuntime *runtimeRT;
     nvinfer1::INetworkDefinition *networkRT; 
+    nvinfer1::IBuilderConfig *configRT;
     
     nvinfer1::ICudaEngine *engineRT;
     nvinfer1::IExecutionContext *contextRT;

src/Int8BatchStream.cpp

@@ -0,0 +1,180 @@
+#include "Int8BatchStream.h"
+
+BatchStream::BatchStream(tk::dnn::dataDim_t dim, int batchSize, int maxBatches, const std::string& fileimglist, const std::string& filelabellist)
+{
+    mBatchSize = batchSize;
+    mMaxBatches = maxBatches;
+    mDims = nvinfer1::DimsNCHW{ dim.n, dim.c, dim.h, dim.w };
+    mHeight = dim.h;
+    mWidth = dim.w;
+    mImageSize = mDims.c()*mDims.h()*mDims.w();
+    mBatch.resize(mBatchSize*mImageSize, 0);
+    mLabels.resize(mBatchSize, 0);
+    mFileBatch.resize(mDims.n()*mImageSize, 0);
+    mFileLabels.resize(mDims.n(), 0);
+    mFileImgList = fileimglist;
+    readInListFile(fileimglist, mListImg);
+    mFileLabelList = filelabellist;
+    readInListFile(filelabellist, mListLabel);
+
+    reset(0);
+}
+
+void BatchStream::reset(int firstBatch)
+{
+    mBatchCount = 0;
+    mFileCount = 0;
+    mFileBatchPos = mDims.n();
+    skip(firstBatch);
+}
+
+bool BatchStream::next()
+{
+    std::cout<<"Next batch: "<<mBatchCount<<" of "<<mMaxBatches<<"\n";
+    if (mBatchCount == mMaxBatches-1)
+        return false;
+
+    for (int csize = 1, batchPos = 0; batchPos < mBatchSize; batchPos += csize, mFileBatchPos += csize)
+    {
+        assert(mFileBatchPos > 0 && mFileBatchPos <= mDims.n());
+        if (mFileBatchPos == mDims.n() && !update())
+            return false;
+
+        csize = std::min(mBatchSize - batchPos, mDims.n() - mFileBatchPos);
+        std::copy_n(getFileBatch() + mFileBatchPos * mImageSize, csize * mImageSize, getBatch() + batchPos * mImageSize);
+        std::copy_n(getFileLabels() + mFileBatchPos, csize, getLabels() + batchPos);
+    }
+    mBatchCount++;
+    return true;
+}
+
+void BatchStream::skip(int skipCount)
+{
+    if (mBatchSize >= mDims.n() && mBatchSize%mDims.n() == 0 && mFileBatchPos == mDims.n())
+    {
+        mFileCount += skipCount * mBatchSize / mDims.n();
+        return;
+    }
+
+    int x = mBatchCount;
+    for (int i = 0; i < skipCount; i++)
+        next();
+    mBatchCount = x;
+}
+
+void BatchStream::readInListFile(const std::string& dataFilePath, std::vector<std::string>& mListIn)
+{
+    // dataFilePath contains the list of image paths
+    int count = 0;
+    FILE* f = fopen(dataFilePath.c_str(), "r");
+    if (!f)
+        FatalError("failed to open " + dataFilePath);
+    
+    char str[512];
+    while (fgets(str, 512, f) != NULL){
+        for (int i = 0; str[i] != '\0'; ++i){
+            if (str[i] == '\n'){
+                str[i] = '\0';
+                break;
+            }
+        }
+        count ++;
+        mListIn.push_back(str);
+        if(count == mMaxBatches)
+            break;
+    }
+    fclose(f);
+}
+
+void BatchStream::readCVimage(std::string inputFileName, std::vector<float>& res, bool fixshape)
+{
+    // unaltered original DsImage
+    cv::Mat m_OrigImage;
+    // letterboxed DsImage given to the network as input
+    cv::Mat m_LetterboxImage;
+    m_OrigImage = cv::imread(inputFileName, cv::IMREAD_COLOR);
+
+    if (!m_OrigImage.data || m_OrigImage.cols <= 0 || m_OrigImage.rows <= 0)
+        FatalError("Unable to open " + inputFileName);
+
+    int m_Height = m_OrigImage.rows;
+    int m_Width = m_OrigImage.cols;
+    if(fixshape){
+        m_Height = mHeight;
+        m_Width = mWidth;
+    }
+    std::cout<<"image is "<<inputFileName<<": "<<m_Height<<" * "<<m_Width<<std::endl;
+    // resize the DsImage with scale
+    float dim = std::max(m_Height, m_Width);
+    int resizeH = ((m_Height / dim) * m_Height);
+    int resizeW = ((m_Width / dim) * m_Width);
+    float m_ScalingFactor = static_cast<float>(resizeH) / static_cast<float>(m_Height);
+
+    // Additional checks for images with non even dims
+    if ((m_Width - resizeW) % 2) resizeW--;
+    if ((m_Height - resizeH) % 2) resizeH--;
+    assert((m_Width - resizeW) % 2 == 0);
+    assert((m_Height - resizeH) % 2 == 0);
+
+    int m_XOffset = (m_Width - resizeW) / 2;
+    int m_YOffset = (m_Height - resizeH) / 2;
+
+    assert(2 * m_XOffset + resizeW == m_Width);
+    assert(2 * m_YOffset + resizeH == m_Height);
+
+    // resizing
+    cv::resize(m_OrigImage, m_LetterboxImage, cv::Size(resizeW, resizeH), 0, 0, cv::INTER_CUBIC);
+    // letterboxing
+    cv::copyMakeBorder(m_LetterboxImage, m_LetterboxImage, m_YOffset, m_YOffset, m_XOffset,
+                    m_XOffset, cv::BORDER_CONSTANT, cv::Scalar(128, 128, 128));
+    m_LetterboxImage.convertTo(m_LetterboxImage, CV_32FC3, 1 / 255.0);
+    // converting to RGB and NCHW format
+    m_LetterboxImage = cv::dnn::blobFromImage(m_LetterboxImage);
+    res.assign(m_LetterboxImage.begin<float>(), m_LetterboxImage.end<float>());
+}
+
+void BatchStream::readLabels(std::string inputFileName, std::vector<float>& ris)
+{
+    std::ifstream is(inputFileName.c_str());
+    //read only the first number: the image sub-portion class
+    while (true) {
+        float val;
+        // Read
+		is >> val;
+        // Check
+        if (!is) {
+            break;
+        }
+        // Use
+        // insert the first number and skip all others
+        ris.push_back(val);
+        while( true ){
+            char c;
+            is >> c;
+            if (is.peek() == '\n') //detect "\n"
+                break;
+        }   
+    }
+}
+
+bool BatchStream::update()
+{
+    std::string imgFileName = mListImg[mFileCount];
+    std::string labelFileName = mListLabel[mFileCount];
+    mFileCount++;
+
+    //read image
+    mFileBatch.clear();
+    readCVimage(imgFileName, mFileBatch);
+    // std::transform(
+    //     singleImg_rawData.begin(), singleImg_rawData.end(), mFileBatch.begin(), [](uint8_t val) { return static_cast<float>(val); });
+    
+    //read label
+    mFileLabels.clear();
+    readLabels(labelFileName, mFileLabels);
+    // std::transform(
+    //     singleLabels_rawData.begin(), singleLabels_rawData.end(), mFileLabels.begin(), [](uint8_t val) { return static_cast<float>(val); });
+    
+    mFileBatchPos = 0;
+    return true;
+}

src/Int8Calibrator.cpp

@@ -0,0 +1,51 @@
+#include "Int8Calibrator.h"
+
+Int8EntropyCalibrator::Int8EntropyCalibrator(BatchStream& stream, int firstBatch, 
+                                             const std::string& calibTableFilePath,
+                                             const std::string& inputBlobName,
+                                             bool readCache): 
+    mStream(stream), 
+    mCalibTableFilePath(calibTableFilePath),
+    mInputBlobName(inputBlobName.c_str()),
+    mReadCache(readCache)
+{
+    nvinfer1::DimsNCHW dims = mStream.getDims();
+    mInputCount = mStream.getBatchSize() * dims.c() * dims.h() * dims.w();
+    checkCuda(cudaMalloc(&mDeviceInput, mInputCount * sizeof(float)));
+    mStream.reset(firstBatch);
+    std::cout<<"mCalibTableFilePath\n";
+}
+
+bool Int8EntropyCalibrator::getBatch(void* bindings[], const char* names[], int nbBindings) 
+{
+    if (!mStream.next())
+        return false;
+
+    checkCuda(cudaMemcpy(mDeviceInput, mStream.getBatch(), mInputCount * sizeof(float), cudaMemcpyHostToDevice));
+    assert(!strcmp(names[0], mInputBlobName.c_str()));
+    bindings[0] = mDeviceInput;
+    return true;
+}
+
+const void* Int8EntropyCalibrator::readCalibrationCache(size_t& length) 
+{
+    mCalibrationCache.clear();
+    assert(!mCalibTableFilePath.empty());
+    std::ifstream input(mCalibTableFilePath, std::ios::binary);
+    input >> std::noskipws;
+    input >> std::noskipws;
+    if (mReadCache && input.good())
+        std::copy(std::istream_iterator<char>(input), std::istream_iterator<char>(),
+                  std::back_inserter(mCalibrationCache));
+
+    length = mCalibrationCache.size();
+    return length ? &mCalibrationCache[0] : nullptr;
+}
+
+void Int8EntropyCalibrator::writeCalibrationCache(const void* cache, size_t length)
+{
+    assert(!mCalibTableFilePath.empty());
+    std::ofstream output(mCalibTableFilePath, std::ios::binary);
+    output.write(reinterpret_cast<const char*>(cache), length);
+    output.close();
+}

src/NetworkRT.cpp

@@ -9,7 +9,7 @@
 #include "utils.h"
 #include "NvInfer.h"
 #include "NetworkRT.h"
-// #include "calibrator.h"
+#include "Int8Calibrator.h"
 
 using namespace nvinfer1;
 
@@ -38,8 +38,16 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
     std::cout<<"Int8 support: "<<builderRT->platformHasFastInt8()<<"\n";
     std::cout<<"DLAs: "<<builderRT->getNbDLACores()<<"\n";
     networkRT = builderRT->createNetwork();
-
+    configRT = builderRT->createBuilderConfig();
+    
     if(!fileExist(name)) {
+        // Calibrator life time needs to last until after the engine is built.
+        std::unique_ptr<IInt8EntropyCalibrator> calibrator;
+
+        configRT->setAvgTimingIterations(1);
+        configRT->setMinTimingIterations(1);
+        configRT->setMaxWorkspaceSize(1 << 30);
+        configRT->setFlag(BuilderFlag::kDEBUG);
 
         //input and dataType
         dataDim_t dim = net->layers[0]->input_dim;
@@ -51,6 +59,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
         if(net->fp16 && builderRT->platformHasFastFp16()) {
             dtRT = DataType::kHALF;
             builderRT->setHalf2Mode(true);
+            configRT->setFlag(BuilderFlag::kFP16);
         }
         if(net->dla && builderRT->getNbDLACores() > 0) {
             dtRT = DataType::kHALF;
@@ -59,16 +68,21 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
             builderRT->setDefaultDeviceType(DeviceType::kDLA);
             builderRT->setDLACore(0);
         }
-        // if(net->int8 && builderRT->platformHasFastInt8())
-        // {
-        //     dtRT = DataType::kINT8;
-        //     builderRT->setInt8Mode(true);
-        //     Int8EntropyCalibrator calibrator(1, "../demo/images.txt","../demo/yolov3-calibration.table", 416*416*3, 416, 416);
-        //     builderRT->setInt8Calibrator((nvinfer1::IInt8Calibrator * )&calibrator);
-        //     // builderRT->setStrictTypeConstraints(true);
-        // }
-       
-        //add input layer
+        if(net->int8 && builderRT->platformHasFastInt8()){
+            // dtRT = DataType::kINT8;
+            // builderRT->setInt8Mode(true);
+            configRT->setFlag(BuilderFlag::kINT8);
+            BatchStream calibrationStream(dim, 1, 100,      //TODO: check if 100 images are sufficient to the calibration (or 4951) 
+                                            "/home/xavier/Documents/tkDNN/demo/COCO_val2017/all_images.txt", 
+                                            "/home/xavier/Documents/tkDNN/demo/COCO_val2017/all_labels.txt");
+            std::string modelName = name;
+            calibrator.reset(new Int8EntropyCalibrator(calibrationStream, 1, 
+                                            "./" + modelName.substr(0, modelName.find('.')) + "-calibration.table", 
+                                            "data"));
+            configRT->setInt8Calibrator(calibrator.get());
+        }
+        
+        // add input layer
         ITensor *input = networkRT->addInput("data", DataType::kFLOAT, 
                         DimsCHW{ dim.c, dim.h, dim.w});
         checkNULL(input);
@@ -77,6 +91,10 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
         for(int i=0; i<net->num_layers; i++) {
             Layer *l = net->layers[i];
             ILayer *Ilay = convert_layer(input, l);
+            if(net->int8 && builderRT->platformHasFastInt8())
+            {
+                Ilay->setPrecision(DataType::kINT8);
+            }
             Ilay->setName( (l->getLayerName() + std::to_string(i)).c_str() );
             
             input = Ilay->getOutput(0);
@@ -94,7 +112,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
         networkRT->markOutput(*input);
 
         std::cout<<"Building tensorRT cuda engine...\n";
-        engineRT = builderRT->buildCudaEngine(*networkRT);
+        engineRT = builderRT->buildEngineWithConfig(*networkRT, *configRT);
         if(engineRT == nullptr)
             FatalError("cloud not build cuda engine")
         // we don't need the network any more