Pytorch使用autograd.function自定义op
0x01 简介
Pytorch是利用Tensor与Function来构建计算图的,Function就像计算图中的边,实现Tensor的计算,并输出新的Tensor,因此Function需要有forward和backward的功能,前者用于正常传递输入,后者用于计算梯度.(但直接用nn.Module构建就可以只用写forward,backward调用自动求导计算了,另用Module可以保存参数而Function不能,因此前者多用于写一些需要保存参数的如自定义的层,而后者通常用来写一个操作,如激活函数之类的,偷懒直接nn.Module就行了).
0x02 使用
pytorch官方文档上写的是得用save_for_backward保存下输入,这个保存的数据会在backward的时候通过saved_tensors读取
使用时调用apply方法即可。
或者直接重命名一下
你已经学会了,现在来试试手写一个Relu吧.jpg
import torch
from torch.autograd import Variable
from torch.autograd import Function
class Myop1(Function):
@staticmethod
def forward(self, input):
self.save_for_backward(input)
output = input.clamp(min=0)
return output
@staticmethod
def backward(self, grad_output):
##
input = self.saved_tensors
grad_input = grad_output.clone()
grad_input[input<0] = 0
return grad_input
看一下Variable与Function的关系
input_1 = Variable(torch.randn(1), requires_grad=True)
print(input_1)
relu2 = Myop1.apply
output_= relu2(input_1)
print(output_.grad_fn)
输出
<torch.autograd.function.Myop1Backward object at 0x00000123A49A29E0>
or封装成函数
def relu(input):
return Myop1.apply(input)