分类：TensorRT5

2020
09-29

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解析器构建EnginedefONNX_build_engine(onnx_file_path):'''通过加载on... 继续阅读 >