zl程序教程

您现在的位置是:首页 >  后端

当前栏目

Pytorch部署方案(一):基于TensorRT(一个C++库)【最成熟&最优化】【①Pytorch->ONNX->TensorRT;②Pytorch->Caffe->TensorRT】

C++amp部署PyTorch 一个 基于 方案 最优化
2023-09-27 14:20:40 时间

一、Pytorch转ONNX的意义

一般来说转ONNX只是一个手段,在之后得到ONNX模型后还需要再将它做转换,比如转换到TensorRT上完成部署,或者有的人多加一步,从ONNX先转换到caffe,再从caffe到tensorRT。原因是Caffe对tensorRT更为友好,这里关于友好的定义后面会谈。

因此在转ONNX工作开展之前,首先必须明确目标后端。ONNX只是一个格式,就和json一样。只要你满足一定的规则,都算是合法的,因此单纯从Pytorch转成一个ONNX文件很简单。但是不同后端设备接受的onnx是不一样的,因此这才是坑的来源。

Pytorch自带的torch.onnx.export转换得到的ONNX,ONNXRuntime需要的ONNX,TensorRT需要的ONNX都是不同的。

真正落地的算法部署,毫无疑问,假如你的target是GPU,采用ONNX+TensorRT应该是目前最成熟、最优化的方案。

假如你的target是一些嵌入式芯片,那么采用MNN也是可以通过onnx轻松实现CPU嵌入式端快速推理的。

既然ONNX和TensorRT这么好,为什么都不用呢?为什么都还在用Python写难看的推理的代码呢?原因也很简单:

  • 入门太难,C++一般人玩不来;
  • 既要懂模型的每一层输入输出,又要懂TensorRT的API,至少要很熟悉。

二、TensorRT

一般的深度学习项目,训练时为了加快速度&#x