perseusdg
72 lines
1.9 KiB
C++
72 lines
1.9 KiB
C++
#ifndef INT8BATCHSTREAM_H
|
|
#define INT8BATCHSTREAM_H
|
|
|
|
#include <vector>
|
|
#include <assert.h>
|
|
#include <algorithm>
|
|
#include <iterator>
|
|
#include <stdint.h>
|
|
#include <iostream>
|
|
#include <string>
|
|
#include <fstream>
|
|
#include <iomanip>
|
|
#include <signal.h>
|
|
#include <stdlib.h>
|
|
#ifdef __linux__
|
|
#include <unistd.h>
|
|
#endif
|
|
|
|
#include <mutex>
|
|
|
|
#include "NvInfer.h"
|
|
#include "utils.h"
|
|
#include "tkdnn.h"
|
|
|
|
/*
|
|
* BatchStream implements the stream for the INT8 calibrator.
|
|
* It reads the two files .txt with the list of image file names
|
|
* and the list of label file names.
|
|
* It then iterates on images and labels.
|
|
*/
|
|
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::Dims4 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 };
|
|
int mFileBatchPos{ 0 };
|
|
int mImageSize{ 0 };
|
|
|
|
nvinfer1::Dims4 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
|