贤者之路,Tensorrt的int8 calibration创建

【INT8 Calibration】

无论哪块GPU计算板卡都能够支持Int8的加速,但需要事先生成 calibration文件,下面是.h .cpp分别继承IInt8EntropyCalibrator类,用于创建Int8 table文件

#pragma once
#include "NvInfer.h"
#include "NvCaffeParser.h"
#include 
#include 
#include 
#include "util.h"

using namespace nvinfer1;
using namespace nvcaffeparser1;


class YoloInt8 : public IInt8EntropyCalibrator
{
public:
	YoloInt8() {}
	virtual ~YoloInt8() {};
	
    // ======================== 初始化实现 ==============================//

	int init(const char* inDir, const char* outDir, int batch, int maxBatch);
	
    // ======================== 基类函数重写实现 ======================= // 
	int destroy();

	int getBatchSize() const override;

	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;
	
	int readStringFromTxt(std::string dir, std::vector &files);
	

    // ======================= 图像处理的辅助函数 ==================== //
    int convMat2Buf_32FC3(float* h_buff, cv::Mat& mat);
	
	int convMat2Buf_8UC3(uchar* h_buff, cv::Mat& srcMat);	


public:
	
	std::vector _fileContainer;

	// 基础参数
	int _width = 608;
	int _height = 416;
	int _channel = 3; 
	int MAXBATCH;
	int _fullsize =  _width * _height * _channel;//  +  net_width * net_height / 2;
	int _singlesize =  _width * _height;//  +  net_width * net_height / 2;

	// getbatch处理时的cpu内存,gpu内存地址
	void* _d_buff;
	void* _h_buff;
	
        // batch变量
	int _batch = 0;
	int _batchIdx = 0;
	std::string _outDirInt8; //filename to save the generated calibration table
	
};

【关键】 

关键是要实现getBatchSize()函数和getBatch()

bool YoloInt8::getBatch(void* bindings[], const char* names[], int nbBindings)
{
	
	if (_batch == 0)
	{
		printf("【error】:please init jlYoloInt8 first\n");
		return false;
	}
	
	if (_batchIdx >= MAXBATCH)
		return false;
	
	float* buffOffset = nullptr;
	
	cv::Mat src;
	cv::Mat tmp(cv::Size(608, 416), CV_32FC3);
	cv::Mat res;
	for (int nPic = 0; nPic < _batch; nPic++) 
	{
		int random = rand() % _fileContainer.size();

		src= cv::imread(_fileContainer[random]);

		cv::resize(src, tmp, cv::Size(608, 416), 0, 0, CV_INTER_CUBIC);
		
		tmp.convertTo(res,CV_32FC3);
		
		buffOffset = (float*)_h_buff + _fullsize * nPic;	
		convMat2Buf_32FC3(buffOffset, finalMat);	
	}
	_batchIdx += _batch;

	cudaMemcpy(_d_buff, _h_buff, _fullsize * _batch * sizeof(float),     cudaMemcpyHostToDevice);
	assert(!strcmp(names[0], JL_YOLO_INPUT_BLOB_NAME));
	bindings[0] = _d_buff;
	return true;
}

 

 

 

你可能感兴趣的:(贤者之路,Tensorrt的int8 calibration创建)