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)