自定义pytorch算子并导出可推理的onnx文件

何为自定义算子

自定义算子(custom operaotors)其实在pytorch内并没有一个特别明确的定义，使用Pytorch提供的自定义算子指南内的定义:

An operator is glue code for the PyTorch runtime that tells it about the computation. A single operator can be associated with multiple kernels (for example, torch.add has a kernel for CPU and a kernel for CUDA). The glue code is necessary to get PyTorch subsystems (like torch.compile and torch.autograd) to compose with the computation.

自定义算子是一种胶水代码块用来告知pytorch runtime如何进行Tensor计算。自定义算子是能够被pytorch原生的子系统(例如torch.complie和torch.autograd等)支持并可以直接在Pytorch框架内使用的。最重要的一点，自定义算子可以拥有不同的kernel来实现。自定义算子将会被pytorch当做一个黑盒进行不透明调用(opaque call)。

为什么在这里强调定义，是因为之前看过一些博客和开源算子库(例如这篇博客和bevpool_v2的官方实现库)，他们将继承于torch.autograd.Function的自定义函数也称作”自定义算子”。实际上这和Pytorch官方给出的自定义算子有所差距。可以看到官方文档内的torch.autograd.Function关联页面介绍了什么时候我们应该继承torch.autograd：

In general, implement a custom function if you want to perform computations in your model that are not differentiable or rely on non-PyTorch libraries (e.g., NumPy), but still wish for your operation to chain with other ops and work with the autograd engine.

如果我们想定义一个正常来说不可导的函数(例如伪量化算子和稀疏掩码矩阵的计算，可以参考我博客的另外一篇文章)或者函数需要依赖于某些Pytorch之外的库来进行forward计算(或backward)，这时候我们可以考虑继承autograd.Function来实现一个自定义的Python内的Pytroch函数。这和前面提到的自定义算子最大的区别在于，该函数是不完全被Pytorch原生子系统所支持的，是单纯在Python侧实现的。当我们需要使用诸如Pytorch的TorchScript/Dynamo等功能或使用C++ API时，就必须使用真正的自定义算子。

当然直接继承autograd也有一些好处，如果只是单纯的想在Python运行时(不管是训练还是推理)调用外部函数(例如cuda)并且不需要关心使用Pytorch的高级功能，那autograd也可以完成上述使命并且写起来也更简单。之前提到的bevpool_v2在智能驾驶的视觉模型内经常会使用到，继承autograd可以让bevpool快速集成入torch框架并参与训练与python侧的推理，但在导出onnx时可能会遇到不少问题。如何继承autograd可以直接看上面提到的博客，本文更偏向于介绍如何使用torch.library.define和@custom_opapi定义算子。

如何编写一个自定义算子

假设我们有一个新的算子需求，并且我们有它的cuda实现(一些.cu文件)或者我们可以调用cuBLAS内的函数来实现它，我们可以使用Python或C++ API进行算子定义。当然，使用cpp进行自定义算子开发相比python而言肯定会复杂很多(主要是cpp的引用以及编译问题)，这里我们主要探讨使用Python API进行定义的方法。就以bev_pool_v2为例，先定义它的描述符：

@torch.library.custom_op("bev::bev_pool_v2", mutates_args=())
def bev_pool_v2(g, depth, feat, ranks_depth, ranks_fea)