Pytorch通过保存为ONNX模型转TensorRT5的实现
1Pytorch以ONNX方式保存模型
defsaveONNX(model,filepath): ''' 保存ONNX模型 :parammodel:神经网络模型 :paramfilepath:文件保存路径 ''' #神经网络输入数据类型 dummy_input=torch.randn(self.config.BATCH_SIZE,1,28,28,device='cuda') torch.onnx.export(model,dummy_input,filepath,verbose=True)
2利用TensorRT5中ONNX解析器构建Engine
defONNX_build_engine(onnx_file_path):
'''
通过加载onnx文件，构建engine
:paramonnx_file_path:onnx文件路径
:return:engine
'''
#打印日志
G_LOGGER=trt.Logger(trt.Logger.WARNING)
withtrt.Builder(G_LOGGER)asbuilder,builder.create_network()asnetwork,trt.OnnxParser(network,G_LOGGER)asparser:
builder.max_batch_size=100
builder.max_workspace_size=1<<20
print('LoadingONNXfilefrompath{}...'.format(onnx_file_path))
withopen(onnx_file_path,'rb')asmodel:
print('BeginningONNXfileparsing')
parser.parse(model.read())
print('CompletedparsingofONNXfile')
print('Buildinganenginefromfile{};thismaytakeawhile...'.format(onnx_file_path))
engine=builder.build_cuda_engine(network)
print("CompletedcreatingEngine")
#保存计划文件
#withopen(engine_file_path,"wb")asf:
#f.write(engine.serialize())
returnengine
3构建TensorRT运行引擎进行预测
defloadONNX2TensorRT(filepath):
'''
通过onnx文件，构建TensorRT运行引擎
:paramfilepath:onnx文件路径
'''
#计算开始时间
Start=time()
engine=self.ONNX_build_engine(filepath)
#读取测试集
datas=DataLoaders()
test_loader=datas.testDataLoader()
img,target=next(iter(test_loader))
img=img.numpy()
target=target.numpy()
img=img.ravel()
context=engine.create_execution_context()
output=np.empty((100,10),dtype=np.float32)
#分配内存
d_input=cuda.mem_alloc(1*img.size*img.dtype.itemsize)
d_output=cuda.mem_alloc(1*output.size*output.dtype.itemsize)
bindings=[int(d_input),int(d_output)]
#pycuda操作缓冲区
stream=cuda.Stream()
#将输入数据放入device
cuda.memcpy_htod_async(d_input,img,stream)
#执行模型
context.execute_async(100,bindings,stream.handle,None)
#将预测结果从从缓冲区取出
cuda.memcpy_dtoh_async(output,d_output,stream)
#线程同步
stream.synchronize()
print("TestCase:"+str(target))
print("Prediction:"+str(np.argmax(output,axis=1)))
print("tensorrttime:",time()-Start)
delcontext
delengine
补充知识：Pytorch/Caffe可以先转换为ONNX，再转换为TensorRT
近来工作，试图把Pytorch用TensorRT运行。折腾了半天，没有完成。github中的转换代码，只能处理pytorch0.2.0的功能（也明确表示不维护了）。和同事一起处理了很多例外，还是没有通过。吾以为，实际上即使勉强过了，能不能跑也是问题。
后来有高手建议，先转换为ONNX，再转换为TensorRT。这个思路基本可行。
是不是这样就万事大吉？当然不是，还是有严重问题要解决的。这只是个思路。
以上这篇Pytorch通过保存为ONNX模型转TensorRT5的实现就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持毛票票。