pytorch官方提供了注册自己的自定义算子的接口,不需要像native_functions.yaml那样每次改源码,再编译
在pytorch仓库下找一个地方,创建新的文件夹
我建议在/pytorch/aten/src/ATen/core/op_registration下创建
创建的文件夹假设叫做myrelu
在myrelu下继续创建myrelu.cpp CMakeLists.txt 和 build文件夹
myrelu.cpp:
#include <torch/script.h>
torch::Tensor myrelu(torch::Tensor self){
return self * (self > 0);
}
TORCH_LIBRARY(myop, m){
m