在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;
}