YOLOv5 tensorRT C++代码详解之engine的读取

在tensorRT中，engine模型是对yolov5序列化后的结果。在推理的时候需要进行反序列化才能进行推理。那么第一步就是读取engine文件。

ifstream介绍

engine文件的读取需要用到ifstream，需要导入头文件#include <iostream>，#include <fstream>.主要提供了三个类：

ifstream – 从已有的文件读入
ofstream – 向文件写内容
fstream - 打开文件供读写

文件打开模式：
ios::in 只读
ios::out 只写
ios::app 从文件末尾开始写，防止丢失文件中原来就有的内容
ios::binary 二进制模式
ios::nocreate 打开一个文件时，如果文件不存在，不创建文件
ios::noreplace 打开一个文件时，如果文件不存在，创建该文件
ios::trunc 打开一个文件，然后清空内容
ios::ate 打开一个文件时，将位置移动到文件尾

文件指针位置在c++中的用法：
ios::beg 文件头
ios::end 文件尾
ios::cur 当前位置

engine文件的读取

读取文件名为"yolov5s.engine"的trt文件，以二进制进行读取。

std::string engine_name = "yolov5s.engine";
std::ifstream file(engine_name, std::ios::binary);

可以判断是否读取成功：

if (!file.good())
	{
		std::cerr<<"read" <<engine_name <<"error!"<<std::endl;
		return -1;
	}

如果可以成功读取，我们将获取完整的数据内容：

读取完整的文件内容，并利用file.tellg()返回文件指针位置,可以获得文件大小；

size_t size = 0;
file.seekg(0, file.end);
size = file.tellg();

size =  20355643 

定义空网络

char* trtModel = nullptr; //trtmodel指针初始化
trtModel = new char[size];
assert(trtModel);

可以看到此时定义的trtModel是一个char类型的指针，并初始化为空指针。其实这一步就是为了在内存中创建一段空间用来存储engine模型。

将engine写入创建的trtModel

    //将engine读入定义trt的内存空间中
	file.read(trtModel, size);
	file.close();

将engine文件放入前面定义的trtModel中。

完整代码

#include <stdio.h>
#include <string>
#include <fstream>
#include <iostream>
#include <assert.h>

int main()
{	
	std::string engine_name = "yolov5s.engine";
	std::ifstream file(engine_name, std::ios::binary);
	if (!file.good())
	{
		std::cerr<<"read " <<engine_name <<" error!"<<std::endl;
		return -1;
	}
	/*std::string line;
	while (std::getline(file, line))
	{
		std::cout << line << std::endl;
	}*/
	
	//如果可以成功读取文件名
	size_t size = 0;
	file.seekg(0, file.end);
	size = file.tellg();
	//定义trt网络
	char* trtModel = nullptr; //trtmodel指针初始化
	trtModel = new char[size];
	assert(trtModel);

	//将engine读入定义trt的内存空间中
	file.read(trtModel, size);
	file.close();
	
	return 0;
}