tensorRT自定义算子

tensorflow建自定义算子

tensorflow建自定义算子
参考：https://blog.csdn.net/xiangxianghehe/article/details/81002227
Compute应该是固定的名字
或者这个：https://www.tensorflow.org/guide/create_op?hl=zh-cn
请先创建一个名为 zero_out.cc 的文件
要创建其中一个内核，请先创建一个扩展 OpKernel 并重写 OpKernel 方法的类。
实现内核后，您需要将其注册到 TensorFlow 系统。在注册中，您要指定此内核将在哪些不同约束下运行。
https://tensorflow.juejin.im/extend/adding_an_op.html

s = decoded_values->allocate(p, TensorShape({p_num}), &p_values);
https://blog.csdn.net/sinat_37532065/article/details/92085177

设置输出维度的参考在这里：https://stackoverflow.com/questions/56573357/tensorflow-how-to-define-shape-of-output-in-setshapefntensorflowshape-i
 

pytorch建自定义算子

torch.ops.load_library("/Users/admin//pytorch/aten/src/ATen/core/op_registration/myrelu/build/libmyrelu.so")
https://www.pythonheidong.com/blog/article/484431/b15d025d5979d5db8a67/

PyTorch自定义CUDA算子

main.py，这是python入口，也就是你平时写模型的地方。
add2.cpp，这是torch和CUDA连接的地方，将CUDA程序封装成了python可以调用的库。
add2.h，CUDA函数声明。
add2.cu，CUDA函数实现。

Torch C++封装：
#include <torch/extension.h>

torch_launch_add2函数传入的是C++版本的torch tensor，然后转换成C++指针数组，调用CUDA函数launch_add2来执行核函数。

这里用pybind11来对torch_launch_add2函数进行封装，然后用cmake编译就可以产生python可以调用的.so库。但是我们这里不直接手动cmake编译

Python调用：
from torch.utils.cpp_extension import load

torch.utils.cpp_extension.load函数就是用来自动编译上面的几个cpp和cu文件的。

总结一下，主要分为三个模块：
先编写CUDA算子和对应的调用函数。
然后编写torch cpp函数建立PyTorch和CUDA之间的联系，用pybind11封装。
最后用PyTorch的cpp扩展库进行编译和调用。
https://godweiyang.com/2021/03/18/torch-cpp-cuda/
https://zhuanlan.zhihu.com/p/358220419

自定义 C ++和 CUDA 扩展

C ++扩展有两种形式：它们可以使用setuptools提前构建，也可以通过torch.utils.cpp_extension.load()适时构建。

通过编写一个setup.py脚本来构建 C ++扩展，该脚本使用 setuptools 编译我们的 C ++代码。

将支持 CUDA 的 op 与 PyTorch 集成非常简单。编写setup.py脚本
https://pytorch.apachecn.org/docs/1.4/42.html

参考EDVR 的代码

pytorch自定义算子转onnx

def my_group_norm(g, input, num_groups, scale, bias, eps):
    return g.op("ygao::my_group_norm", input, num_groups, scale, bias, epsilon_f=eps)
register_custom_op_symbolic('ygao::custom_group_norm', my_group_norm, 11)
https://blog.csdn.net/Artyze/article/details/107642358

Pytorch、ONNX以及ONNX-TensorRT5.1 customOP注册（将自定义算子注册到onnx ！！！）
https://blog.csdn.net/tfu259/article/details/104019226

How to export Pytorch model with custom op to ONNX and run it in ONNX Runtime
https://github.com/onnx/tutorials/tree/master/PyTorchCustomOperator

torch.onnx
https://pytorch.org/docs/stable/onnx.html
参数说明
https://blog.csdn.net/QFJIZHI/article/details/105245292

Pytorch底层算子扩展
https://zhuanlan.zhihu.com/p/349560723