diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index df050b7..b891377 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -10,6 +10,7 @@ using namespace nvinfer1; #include "pluginsRT/ActivationLeakyRT.cpp" #include "pluginsRT/ReorgRT.cpp" #include "pluginsRT/RegionRT.cpp" +#include "pluginsRT/Int8Calibrator.cpp" // Logger for info/warning/errors class Logger : public ILogger { @@ -62,6 +63,11 @@ NetworkRT::NetworkRT(Network *net, const char *name) { builderRT->setMaxBatchSize(1); builderRT->setMaxWorkspaceSize(1 << 20); + BatchStream bstream({32,dim.c, dim.h, dim.w}, 32, 1); + Int8EntropyCalibrator calib(bstream, 0, false); + builderRT->setInt8Mode(true); + builderRT->setInt8Calibrator(&calib); + std::cout<<"Building tensorRT cuda engine...\n"; engineRT = builderRT->buildCudaEngine(*networkRT); // we don't need the network any more diff --git a/src/pluginsRT/Int8Calibrator.cpp b/src/pluginsRT/Int8Calibrator.cpp new file mode 100644 index 0000000..60def25 --- /dev/null +++ b/src/pluginsRT/Int8Calibrator.cpp @@ -0,0 +1,166 @@ +#include +#include +#include +#include + +#include "NvInfer.h" + +class BatchStream +{ +public: + BatchStream(tkDNN::dataDim_t dim, int batchSize, int maxBatches) + { + mBatchSize = batchSize; + mMaxBatches = maxBatches; + mDims = nvinfer1::DimsNCHW{ dim.n, dim.c, dim.h, 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); + reset(0); + } + + void reset(int firstBatch) + { + mBatchCount = 0; + mFileCount = 0; + mFileBatchPos = mDims.n(); + skip(firstBatch); + } + + bool next() + { + std::cout<<"Next batch: "< 0 && mFileBatchPos <= mDims.n()); + if (mFileBatchPos == mDims.n() && !update()) + return false; + + // copy the smaller of: elements left to fulfill the request, or elements left in the file buffer. + 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 skip(int skipCount) + { + if (mBatchSize >= mDims.n() && mBatchSize%mDims.n() == 0 && mFileBatchPos == mDims.n()) + { + mFileCount += skipCount * mBatchSize / mDims.n(); + std::cout< mBatch; + std::vector mLabels; + std::vector mFileBatch; + std::vector mFileLabels; +}; + + + +class Int8EntropyCalibrator : public IInt8EntropyCalibrator +{ +public: + Int8EntropyCalibrator(BatchStream& stream, int firstBatch, bool readCache = true) + : mStream(stream), mReadCache(readCache) + { + DimsNCHW dims = mStream.getDims(); + mInputCount = mStream.getBatchSize() * dims.c() * dims.h() * dims.w(); + checkCuda(cudaMalloc(&mDeviceInput, mInputCount * sizeof(float))); + mStream.reset(firstBatch); + } + + virtual ~Int8EntropyCalibrator() + { + checkCuda(cudaFree(mDeviceInput)); + } + + int getBatchSize() const override { return mStream.getBatchSize(); } + + bool getBatch(void* bindings[], const char* names[], int nbBindings) override + { + std::cout<<"CALIB request batch\n"; + if (!mStream.next()) + return false; + + checkCuda(cudaMemcpy(mDeviceInput, mStream.getBatch(), mInputCount * sizeof(float), cudaMemcpyHostToDevice)); + bindings[0] = mDeviceInput; + return true; + } + + const void* readCalibrationCache(size_t& length) override + { + mCalibrationCache.clear(); + std::ifstream input("table.calib", std::ios::binary); + input >> std::noskipws; + if (mReadCache && input.good()) + std::copy(std::istream_iterator(input), std::istream_iterator(), std::back_inserter(mCalibrationCache)); + + length = mCalibrationCache.size(); + return length ? &mCalibrationCache[0] : nullptr; + } + + void writeCalibrationCache(const void* cache, size_t length) override + { + std::ofstream output("table.calib", std::ios::binary); + output.write(reinterpret_cast(cache), length); + } + +private: + BatchStream mStream; + bool mReadCache{ true }; + + size_t mInputCount; + void* mDeviceInput{ nullptr }; + std::vector mCalibrationCache; +}; \ No newline at end of file diff --git a/tests/yolo/yolo.cpp b/tests/yolo/yolo.cpp index 50aec86..7c08a32 100644 --- a/tests/yolo/yolo.cpp +++ b/tests/yolo/yolo.cpp @@ -28,7 +28,7 @@ const char *c30_bin = "../tests/yolo/layers/c30.bin"; const char *g31_bin = "../tests/yolo/layers/g31.bin"; const char *output_bin = "../tests/yolo/layers/output.bin"; -int main() { +int main(int argc, char *argv[]) { // Network layout tkDNN::dataDim_t dim(1, 3, 608, 608, 1); @@ -108,7 +108,7 @@ int main() { net.print(); //convert network to tensorRT - tkDNN::NetworkRT netRT(&net, "yolo.rt"); + tkDNN::NetworkRT netRT(&net, argv[1]); dnnType *out_data, *out_data2; // cudnn output, tensorRT output