TensorRT部署(图像分类)之engine生成及反序列化推理(第二讲)

1)日志文件类创建

class TRTLogger : public nvinfer1::ILogger {
public:
	virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override {
		if (severity <= Severity::kINFO) {
			if (severity == Severity::kWARNING) {
				printf("\033[33m%s: %s\033[0m\n", msg);
			}
			else if (severity <= Severity::kERROR) {
				printf("\033[31m%s: %s\033[0m\n", msg);
			}
			else {
				printf("%s: %s\n", msg);
			}
		}
	}
} logger;

日志文件主要用于打印警告错误信息,方便调试

bool build_model() {
	// 这是基本需要的组件
	auto builder = nvinfer1::createInferBuilder(logger);
	auto config = builder->createBuilderConfig();
	auto network = builder->createNetworkV2(1);

	// 通过onnxparser解析器解析的结果会填充到network中,类似addConv的方式添加进去
	auto parser = nvonnxparser::createParser(*network, logger);
	if (!parser->parseFromFile("classifier.onnx", 1)) {
		printf("Failed to parse classifier.onnx\n");

		// 注意这里的几个指针还没有释放,是有内存泄漏的,后面考虑更优雅的解决
		return false;
	}

	int maxBatchSize = 10;
	printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f);
	config->setMaxWorkspaceSize(1 << 28);

	// 如果模型有多个输入,则必须多个profile
	auto profile = builder->createOptimizationProfile();
	auto input_tensor = network->getInput(0);
	auto input_dims = input_tensor->getDimensions();

	// 配置最小、最优、最大范围
	input_dims.d[0] = 1;
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMIN, input_dims);
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kOPT, input_dims);
	input_dims.d[0] = maxBatchSize;
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);
	config->addOptimizationProfile(profile);

	auto engine = builder->buildEngineWithConfig(*network, *config);
	if (engine == nullptr) {
		printf("Build engine failed.\n");
		return false;
	}

	// 将模型序列化,并储存为文件
	auto model_data = engine->serialize();
	FILE* f;
	fopen_s(&f, "engine.trtmodel", "wb");
	fwrite(model_data->data(), 1, model_data->size(), f);
	fclose(f);

	// 卸载顺序按照构建顺序倒序
	printf("Done.\n");
	return true;
}

2)推理引擎主要包括配置文件的设置,onnx文件的解析,生成序列化文件,方便下次反序列化,毕竟解析onnx文件生成推理引擎挺浪费时间

void inference() {

	TRTLogger logger;
	auto engine_data = load_file("engine.trtmodel");
	auto runtime = nvinfer1::createInferRuntime(logger);
	auto engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());
	if (engine == nullptr) {
		printf("Deserialize cuda engine failed.\n");
		runtime->destroy();
		return;
	}

	cudaStream_t stream = nullptr;
	cudaStreamCreate(&stream);
	auto execution_context = engine->createExecutionContext();

	int input_batch = 1;
	int input_channel = 3;
	int input_height = 224;
	int input_width = 224;
	int input_numel = input_batch * input_channel * input_height * input_width;
	float* input_data_host = nullptr;
	float* input_data_device = nullptr;
	cudaMallocHost(&input_data_host, input_numel * sizeof(float));
	cudaMalloc(&input_data_device, input_numel * sizeof(float));

	///
	// image to float
	auto image = cv::imread("input_image.jpg");
	float mean[] = { 0.406, 0.456, 0.485 };
	float std[] = { 0.225, 0.224, 0.229 };

	// 对应于pytorch的代码部分
	cv::resize(image, image, cv::Size(input_width, input_height));
	int image_area = image.cols * image.rows;
	unsigned char* pimage = image.data;
	float* phost_b = input_data_host + image_area * 0;
	float* phost_g = input_data_host + image_area * 1;
	float* phost_r = input_data_host + image_area * 2;
	for (int i = 0; i < image_area; ++i, pimage += 3) {
		// 注意这里的顺序rgb调换了
		*phost_r++ = (pimage[0] / 255.0f - mean[0]) / std[0];
		*phost_g++ = (pimage[1] / 255.0f - mean[1]) / std[1];
		*phost_b++ = (pimage[2] / 255.0f - mean[2]) / std[2];
	}
	///
	cudaMemcpyAsync(input_data_device, input_data_host, input_numel * sizeof(float), cudaMemcpyHostToDevice, stream);

	// 3x3输入,对应3x3输出
	const int num_classes = 1000;
	float output_data_host[num_classes];
	float* output_data_device = nullptr;
	cudaMalloc(&output_data_device, sizeof(output_data_host));

	// 明确当前推理时,使用的数据输入大小
	auto input_dims = execution_context->getBindingDimensions(0);
	input_dims.d[0] = input_batch;

	// 设置当前推理时,input大小
	execution_context->setBindingDimensions(0, input_dims);
	float* bindings[] = { input_data_device, output_data_device };
	bool success = execution_context->enqueueV2((void**)bindings, stream, nullptr);
	cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream);
	cudaStreamSynchronize(stream);

	float* prob = output_data_host;
	int predict_label = std::max_element(prob, prob + num_classes) - prob;  // 确定预测类别的下标
	auto labels = load_labels("labels.txt");
	auto predict_name = labels[predict_label];
	float confidence = prob[predict_label];    // 获得预测值的置信度
	printf("Predict: %s, confidence = %f, label = %d\n", predict_name.c_str(), confidence, predict_label);

	cudaStreamDestroy(stream);
	cudaFreeHost(input_data_host);
	cudaFree(input_data_device);
	cudaFree(output_data_device);
}

3)推理主要包括反序列化文件生成推理引擎;图像预处理;将数据拷贝到GPU完成并行计算,再拷贝到CPU,拿到推理结果

### TensorRT Engine 文件反序列化教程 TensorRT 提供了一种高效的机制用于加载和运行已经生成的序列化模型(即 `.engine` 文件)。以下是关于如何对 TensorRTEngine 文件进行反序列化的详细说明以及示例代码。 #### 1. 反序列化的核心概念 在 TensorRT 中,序列化后的引擎文件可以通过 `nvinfer1::IRuntime` 类中的方法重新加载到内存中。此过程涉及创建一个运行时实例并调用其 `deserializeCudaEngine` 方法[^1]。 #### 2. 示例代码实现 以下是一个完整的 Python 实现示例,展示如何从磁盘读取已保存的 `.engine` 文件并将其反序列化为可执行的推理引擎: ```python import tensorrt as trt def load_engine(engine_file_path): # 初始化 TensorRT 运行时对象 runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING)) with open(engine_file_path, "rb") as f: engine_data = f.read() # 将 .engine 文件内容读入内存 # 调用 deserialize_cuda_engine 函数完成反序列化 engine = runtime.deserialize_cuda_engine(engine_data) return engine if __name__ == "__main__": engine_file_path = "./model.engine" engine = load_engine(engine_file_path) if engine is not None: print("Engine loaded successfully!") else: print("Failed to load the engine.") ``` 上述代码展示了通过 `Runtime` 对象加载 `.engine` 文件的过程,并返回了一个可以用来执行推理任务的引擎实例。 #### 3. 动态维度支持下的反序列化处理 当涉及到动态输入尺寸的情况时,在反序列化之后还需要额外配置上下文 (Context),以便指定实际使用的批量大小或其他参数。具体做法如下所示: ```cpp #include <NvInfer.h> #include <fstream> using namespace nvinfer1; ICudaEngine* loadEngine(const char* filePath){ std::ifstream file(filePath,std::ios::binary); if(!file.good()){ throw std::runtime_error("Could not find engine file!"); } file.seekg(0,std::ios::end); size_t size=file.tellg(); file.seekg(0); char *trtModelStream=new char[size]; file.read(trtModelStream,size); IRuntime* runtime= createInferRuntime(gLogger); ICudaEngine* engine=runtime->deserializeCudaEngine(trtModelStream,size,nullptr); delete[] trtModelStream; return engine; } void createContextWithDynamicShapes(ICudaEngine& engine,int batchSize){ IExecutionContext* context=engine.createExecutionContext(); bool success=context->setBindingDimensions( 0,//假设第一个绑定项是输入数据 Dims4(batchSize,-1,-1,-1)//根据实际情况调整形状定义 ); } ``` 以上 C++ 版本同样实现了类似的逻辑,但增加了针对动态批次的支持功能[^2]。 #### 注意事项 - 序列化与反序列化过程中需确保版本兼容性。 - 如果存在多个输入/输出节点,则应逐一为其分配缓冲区地址。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

CV_er

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值