目錄

• 保存序列化的結(jié)果
• 讀取序列化的結(jié)果

? ? ? ?我們使用TensorRT轉(zhuǎn)化Caffe模型時，每次都要使用如下代碼將模型序列化，之后再進行反序列化，才可以進行inference。但是build engine的過程是十分耗時的，我們可以將序列化的結(jié)果保存至本地，以后直接讀取本地的序列化結(jié)果即可，無需build engine，這樣可以節(jié)省很多時間。

// 此函數(shù)用于得到pluginFactory和gie_model_stream void objectDetection::caffeToGIEModel(const std::string& deploy_file,const std::string& model_file,const std::vector<std::string>& outputs,unsigned int max_batch_size,nvcaffeparser1::IPluginFactory* pluginFactory,IHostMemory **gie_model_stream) {...... //函數(shù)中的其他代碼ICudaEngine *engine = builder->buildCudaEngine(*network);...... //函數(shù)中的其他代碼// 得到序列化的模型結(jié)果(*gie_model_stream) = engine->serialize();...... //函數(shù)中的其他代碼 }// 調(diào)用caffeToGIEModel得到_plugin_factory和_gie_model_stream caffeToGIEModel(_deploy_file, _model_file, _outputs, _batch_size, &_plugin_factory, &_gie_model_stream);// 將序列化得到的結(jié)果進行反序列化，以執(zhí)行后續(xù)的inference _engine = _runtime->deserializeCudaEngine(_gie_model_stream->data(), _gie_model_stream->size(), &_plugin_factory);

保存序列化的結(jié)果

// 此函數(shù)用于得到pluginFactory和gie_model_stream void objectDetection::caffeToGIEModel(const std::string& deploy_file,const std::string& model_file,const std::vector<std::string>& outputs,unsigned int max_batch_size,nvcaffeparser1::IPluginFactory* pluginFactory,IHostMemory **gie_model_stream) {...... //函數(shù)中的其他代碼ICudaEngine *engine = builder->buildCudaEngine(*network);...... //函數(shù)中的其他代碼// 得到序列化的模型結(jié)果(*gie_model_stream) = engine->serialize();// 設(shè)置保存文件的名稱為cached_model.binstd::string cache_path = "../cached_model.bin";std::ofstream serialize_output_stream;// 將序列化的模型結(jié)果拷貝至serialize_str字符串std::string serialize_str;serialize_str.resize( (*gie_model_stream)->size() );memcpy((void*)serialize_str.data(), (*gie_model_stream)->data(), (*gie_model_stream)->size());// 將serialize_str字符串的內(nèi)容輸出至cached_model.bin文件serialize_output_stream.open(cache_path);serialize_output_stream << serialize_str;serialize_output_stream.close();...... //函數(shù)中的其他代碼 }

讀取序列化的結(jié)果

// 調(diào)用caffeToGIEModel得到_plugin_factory，將此函數(shù)中和_gie_model_stream相關(guān)的代碼注釋或者刪除 caffeToGIEModel(_deploy_file, _model_file, _outputs, _batch_size, &_plugin_factory, &_gie_model_stream);// 從cached_model.bin文件中讀取序列化的結(jié)果 std::string cached_path = "../cached_model.bin"; std::ifstream fin(cached_path);// 將文件中的內(nèi)容讀取至cached_engine字符串 std::string cached_engine = ""; while (fin.peek() != EOF){ // 使用fin.peek()防止文件讀取時無限循環(huán)std::stringstream buffer;buffer << fin.rdbuf();cached_engine.append(buffer.str()); } fin.close();// 將序列化得到的結(jié)果進行反序列化，以執(zhí)行后續(xù)的inference _engine = _runtime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(), &_plugin_factory);

總結(jié)

以上是生活随笔為你收集整理的TensorRT学习笔记7 - 保存与读取序列化的结果的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。