TensorRT保存序列化的结果

在上一篇博客中分析了tensorRT在加速caffe推断时的主要路程,其中序列化和反序列化是必不可少的。序列化时根据输入网络的prototxt和训练好的模型参数,对网络进行解析和参数保存,主要函数为:

// serialize the engine, then close everything down
	gieModelStream = engine->serialize();

反序列化则是:

// deserialize the engine 
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(gieModelStream->data(), gieModelStream->size(), nullptr);

在tensorRT推断初始化时,序列化的时间较长,特别是较大的网络,影响开发者的进度。从代码中可见,序列化的结果保存在gieModelStream中,然而对于一个已经与训练好的固定网络,我们只需序列化一次,将gieModelStream中的内容的保存到文件中即可,在运行推断反序列化时从直接从文件中读取保存结果,不在执行耗时的序列化操作即可,这无疑会提高调试和运行的效率。

具体做法如下:

//cppy serialize result from gieModelStream to serialize_str,and save it into "serialize_engine_output.txt".
 std::ofstream serialize_output_stream;
 serialize_str.resize(gieModelStream->size());   
 memcpy((void*)serialize_str.data(),gieModelStream->data(),gieModelStream->size());
  serialize_output_stream.open(serialize_engine_output.txt");
serialize_output_stream<

以mnist为例详细过程的代码可以到我的github去下载(https://github.com/haoxurt/tensorRT_save_serialization)

你可能感兴趣的:(tensorRT)