哈喽,大家好,距离上一篇文章又过了半个多月,今天和大家聊聊如何导出pytorch中的自定义op到onnx,同时使用onnx runtime执行。
本篇文章将介绍基于自定义op扩展TorchScript的一些必要步骤,导出op为onnx格式,然后将该op添加到onnx runtime完成模型的inference。尽管Torch提供了大量通用的tensor操作,但是TorchScript是的用户能够使用C++或者CUDA非常方便的实现他们自己的函数并注册成一个新的op。为了将这类自定义的op导出为onnx格式,onnx 提供了自定义op的注册接口,使得用户能够将包含有原生op以及新的自定义onnx op导出为TorchScript op. 一旦op被转换成了onnx格式,用户就可以实现并注册到onnx runtime中并进行模型推理。本文将通过一个例子,详细描述一下这个过程。
步骤
- 添加一个C++实现的自定义op并使用TorchScript进行注册。
- 将自定义op导出为onnx格式,可是已有onnx op的组合或者一个新的自定义op。
- 添加一个自定义operator的实现并住处到ONNX runtime
实现自定义算子