Explicación detallada del código YOLOv5 tensorRT C ++ de la lectura del motor

En tensorRT, el modelo del motor es el resultado de serializar yolov5. Se requiere deserialización durante el razonamiento para realizar el razonamiento. Entonces el primer paso es leer el archivo del motor.

introducción de ifstream

La lectura de archivos del motor requiere el uso de ifstream y los archivos de encabezado deben importarse #include <iostream>, #include <fstream>. Se proporcionan tres clases principales:

ifstream - leer desde un archivo existente
ofstream - escribir contenido en un archivo
fstream - abrir un archivo para lectura y escritura

Modo de apertura de archivos:
ios::in de solo lectura
ios::out de solo escritura
ios::app Comience a escribir desde el final del archivo para evitar perder el contenido original del archivo
ios::binary Modo binario
ios::nocreate Cuándo Al abrir un archivo, si el archivo no existe, no cree el archivo
ios::noreplace Al abrir un archivo, si el archivo no existe, cree el archivo
ios::trunc Abra un archivo y luego borre el contenido
ios:: comido Al abrir un archivo, mueva la posición a la cola del archivo

Uso de la posición del puntero del archivo en C++:
ios::beg file head
ios::end file end
ios::cur posición actual

Lectura de archivos del motor

Lea el archivo trt llamado "yolov5s.engine" en formato binario.

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

Puede determinar si la lectura es exitosa:

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

Si se puede leer correctamente, obtendremos el contenido de datos completo:

Lea el contenido completo del archivo y use file.tellg() para devolver la posición del puntero del archivo para obtener el tamaño del archivo;

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

tamaño = 20355643 

Definir una red vacía

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

Puede ver que el trtModel definido en este momento es un puntero de tipo char y se inicializa como un puntero nulo. De hecho, este paso consiste en crear un espacio en la memoria para almacenar el modelo del motor.

Escribir motor en el trtModel creado

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

Coloque el archivo del motor en el trtModel definido anteriormente.

código completo

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