pytorch自定义算子 native

pytorch自定义算子 native_functions.yaml

在pytorch的文件夹中搜索native_functions.yaml，可以看到，所有pytorch原生的函数都是在这里注册的

在同文件夹下，存在一个README.md，这是官方教程，教我们怎么用native_functions.yaml注册自己自定义的算子

下面我仿照pytorch原生的relu算子，写一个自己的myrelu

首先，找到native_functions.yaml，有几个函数是我们可以参考的

- func: relu(Tensor self) -> Tensoruse_c10_dispatcher: fullvariants: function, methoddispatch:CPU, CUDA: reluMkldnnCPU: mkldnn_reluQuantizedCPU: relu_quantized_cpu

这里除了dispatch其他可以照搬。dispatch主要是针对不同的backend调用不同的kernel，比如GPU上的调用CUDA的kernel。这里展示的是CPU的，在native_functions.yaml加上

- func: myrelu(Tensor self) -> Tensoruse_c10_dispatcher: fullvariants: function, methoddispatch:CPU: myrelu

接下来，找到原生relu的实现，在pytorch/aten/src/ATen/native/Activation.cpp里面
relu的实现调用了同文件下的threshold，threshold的功能是实现 x > thresh ? x : result
模仿写这个调用太复杂了，而且如果在c++段调用函数的话，真正执行的时候每次调用函数都会经过dispatch，会浪费挺多时间的。因此，我参考另外一个激活函数写自己的myrelu

Tensor hardsigmoid(const Tensor& self) {Tensor result;auto iter = TensorIterator::unary_op(result, self);hardsigmoid_stub(iter.device_type(), iter);return iter.output();
}

hardsigmoid是用直线去近似sigmoid函数，它在这个文件中的实现和我们对relu的实现需求很像。这里大致做的事情是，创建一个空的返回对象，然后创建一个TensorIterator，对象类型是单操作数，然后调用hardsigmoid_stub进行具体的实现，最后返回。
在刚刚的activation.cpp中加入

Tensor myrelu(const Tensor& self) {Tensor result;auto iter = TensorIterator::unary_op(result, self);myrelu_stub(iter.device_type(), iter);return iter.output();
}

接下来，我们要知道还需要加上什么使得myrelu_stub也能像hardsigmoid_stub一样可以用。搜索hardsigmoid_stub，发现文件中还存在一个地方

DEFINE_DISPATCH(hardsigmoid_stub);

这段根据字面意思是说，这是一个分发的对象。
加上

DEFINE_DISPATCH(myrelu_stub);

理论上，activation.h也要看一下，activation.cpp和.h是一个整体
在h文件中存在

using hardsigmoid_fn = void(*)(TensorIterator&);
DECLARE_DISPATCH(hardsigmoid_fn, hardsigmoid_stub);

仿照，在对应的地方分别加上

using myrelu_fn = void(*)(TensorIterator&);
DECLARE_DISPACTH(myrelu_fn, myrelu_stub);

这样感觉算是告一段落，那么问题来了，myrelu_stub应该在哪里实现呢
在刚才activation.cpp的文件夹下，还有一个叫做cpu的文件夹，里面存在一个README.md，里面的说明是如何注册一个CPU的算子
在native/cpu/activation.cpp中搜索hard_sigmoid_stub

REGISTER_DISPATCH(hardsigmoid_stub, &hardsigmoid_kernel);

仿照

REGISTER_DISPATCH(myrelu_stub, &myrelu_kernel);

发现了一个新的kernel，叫做hardsigmoid_kernel
搜索hardsigmoid_kernel

void hardsigmoid_kernel(TensorIterator& iter) {AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "hardsigmoid_cpu", [&] {const scalar_t zero(0.0f);const scalar_t three(3.0f);const scalar_t six(6.0f);using Vec = vec256::Vec256<scalar_t>;const Vec kZeroVec(zero);const Vec kThreeVec(three);const Vec kSixVec(six);cpu_kernel_vec(iter,[&](scalar_t self_val) {return std::min(std::max(self_val + three, zero), six) / six;},[&](Vec self_val) {return vec256::minimum(vec256::maximum(self_val + kThreeVec, kZeroVec),kSixVec) / kSixVec;});});
}

仿照写自己的kernel，这用到了vec256指令集，不清楚具体的语法，仿照着用就行
，具体的用法可以看一下relu真实的kernel threshold_kernel的写法

static void threshold_kernel(TensorIterator& iter,Scalar threshold_scalar,Scalar value_scalar) {AT_DISPATCH_ALL_TYPES(iter.dtype(), "threshold_cpu", [&] {using Vec = Vec256<scalar_t>;scalar_t threshold = threshold_scalar.to<scalar_t>();Vec threshold_v = Vec(threshold);scalar_t value = value_scalar.to<scalar_t>();Vec value_v = Vec(value);cpu_kernel_vec(iter,[&](scalar_t x, scalar_t other) -> scalar_t {return x <= threshold ? value : other;},[&](Vec x, Vec other) -> Vec {return Vec::blendv(other, value_v, x <= threshold_v);});});
}

void myrelu_kernel(TensorIterator& iter) {AT_DISPATCH_ALL_TYPES(iter.dtype(), "myrelu_cpu", [&] {const scalar_t zero(0);using Vec = vec256::Vec256<scalar_t>;const Vec kZeroVec(zero);cpu_kernel_vec(iter,[&](scalar_t self_val) {return self_val > zero ? self_val : zero;},[&](Vec self_val) {return Vec::blendv(x,kZeroVec,x<=kZeroVec);});});
}

这里第一个lambda比较显然，后面调用了blendv方法，具体的功能按照猜想是，如果满足第三个式子，则取第二个Vec的值，否则就取第一个Vec的值
写到这里自然会疑惑，过程中有许多对于backward的定义，比如hardsigmoid_backward_kernel
这个反向传播函数在哪里注册呢？
实际上最开始的README中有提到，需要向tools/derivatives.yaml注册
仿照

- name: hardsigmoid(Tensor self) -> Tensorself: hardsigmoid_backward(grad, self)

- name: myrelu(Tensor self) -> Tensorself: myrelu_backward(grad, self)

类似得，myrelu_backward方法同样需要再native_functions.yaml和activation.h activation.cpp 以及cpu文件夹下的activation.cpp中补上对应的代码

实现之后就可以在本地编译了，本地编译的准备工作我在其他的博客中有提到。因为我是用macos，而且不希望编译其他的附加的backend，只是测试一下，所以提供如下的变量设置

export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"} DEBUG=1 USE_CUDA=0 USE_CUDNN=0 USE_FBGEMM=0 USE_NUMPY=0 BUILD_TEST=0 USE_MKLDNN=0 USE_NNPACK=0 USE_QNNPACK=0 USE_DISTRIBUTED=0 BUILD_CAFFE2_OPS=0 MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++python setup.py install

完成之后用下面的python 文件测试一下

import torch
a = torch.randn(5,3,requires_grad=True)
b = torch.randn(5,3,requires_grad=True)
print(a.myrelu())
print(a.relu())
a.relu().backward(b)
print(a.grad)
a.myrelu().backward(b)
print(a.grad/2)

输出：

tensor([[0.0000, 0.0000, 0.0000],[0.0000, 0.9754, 1.0307],[0.0000, 0.0000, 1.4163],[0.0000, 0.0000, 0.0000],[0.3238, 1.4044, 0.0000]], grad_fn=<MyreluBackward>)
tensor([[0.0000, 0.0000, 0.0000],[0.0000, 0.9754, 1.0307],[0.0000, 0.0000, 1.4163],[0.0000, 0.0000, 0.0000],[0.3238, 1.4044, 0.0000]], grad_fn=<ReluBackward0>)tensor([[ 0.0000,  0.0000,  0.0000],[ 0.0000,  0.8448,  0.8919],[ 0.0000,  0.0000, -0.4559],[ 0.0000,  0.0000,  0.0000],[-0.8337, -1.9224,  0.0000]])
tensor([[ 0.0000,  0.0000,  0.0000],[ 0.0000,  0.8448,  0.8919],[ 0.0000,  0.0000, -0.4559],[ 0.0000,  0.0000,  0.0000],[-0.8337, -1.9224,  0.0000]])

可以看出，两者的结果是完全一致的
另外，自定义的反向传播函数也成功地注册了

pytorch自定义算子 native_functions.yaml相关推荐

pytorch之C++实现自定义算子
自定义算子对于输入 x ,其输出为利用C++实现以上算子,总共只要实现两个文件: setup.py 利用python中提供的setuptools模块完成事先编译流程,将写有算子的C++文件,编译成 ...
这可能是关于Pytorch底层算子扩展最详细的总结了！
1.前言一般情况下,pytorch推荐使用python层的前端语言来构建新的算子.因为pytorch在python层的api已经足够丰富,可以构造出很多自定义的算子.但是有时候出于一些其他方面的 ...
pytorch 矩阵相乘_深入浅出PyTorch（算子篇）
Tensor 自从张量(Tensor)计算这个概念出现后,神经网络的算法就可以看作是一系列的张量计算.所谓的张量,它原本是个数学概念,表示各种向量或者数值之间的关系.PyTorch的张量(torch. ...
自定义算子高性能开发
自定义算子高性能开发在计图中,一共有三种方法来开发自定义的算子: 使用元算子进行组合. 使用Code算子开发自定义算子. 使用计图编译器编译自定义的模块和custom op. 其中,元算子开发是最为 ...
java如何给一个链表定义和传值_如何在CUDA中为Transformer编写一个PyTorch自定义层...
如今,深度学习模型处于持续的演进中,它们正变得庞大而复杂.研究者们通常通过组合现有的 TensorFlow 或 PyTorch 操作符来发现新的架构.然而,有时候,我们可能需要通过自定义的操作符来实现 ...
Pytorch自定义数据集
简述 Pytorch自定义数据集方法,应该是用pytorch做算法的最基本的东西. 往往网络上给的demo都是基于torch自带的MNIST的相关类.所以,为了解决使用其他的数据集,在查阅了torch ...
Halcon初学者知识【11】自定义算子和应用实例
1 算子的约定 halcon的自定义算子,包含如下格式: 算子名称 (图象变量-输入 : 图象变量-输出: 控制变量-输入 : 控制变量-输出 ) 要点是: 在HALCON的所有算子(系统算子,自定义 ...
Pytorch自定义Dataset和DataLoader去除不存在和空的数据
Pytorch自定义Dataset和DataLoader去除不存在和空的数据 [源码GitHub地址]:https://github.com/PanJinquan/pytorch-learning-t ...
使用pytorch自定义DataSet，以加载图像数据集为例，实现一些骚操作
使用pytorch自定义DataSet,以加载图像数据集为例,实现一些骚操作总共分为四步构造一个my_dataset类,继承自torch.utils.data.Dataset 重写__getite ...

pytorch自定义算子 native_functions.yaml

pytorch自定义算子 native_functions.yaml

pytorch自定义算子 native_functions.yaml相关推荐

最新文章

热门文章