Pytorch使用autograd.function自定义op

0x01 简介

Pytorch是利用Variable与Function来构建计算图的,Function就像计算图中的边,实现Variable的计算,并输出新的Variable,因此Function需要有forward和backward的功能,前者用于正常传递输入,后者用于计算梯度.(但直接用nn.Module构建就可以只用写forward,backward调用自动求导计算了,另用Module可以保存参数而Function不能,因此前者多用于写一些需要保存参数的如自定义的层,而后者通常用来写一个操作,如激活函数之类的,偷懒直接nn.Module就行了).

0x02 使用

pytorch官方文档上写的是得用save_for_backward保存下输入,这个保存的数据会在backward的时候通过saved_tensors读取

Pytorch使用autograd.function自定义op
使用时调用apply方法即可。
Pytorch使用autograd.function自定义op
或者直接重命名一下
Pytorch使用autograd.function自定义op

你已经学会了,现在来试试手写一个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)
上一篇:新版本浏览器跨站cookie 获取问题


下一篇:matlab递归神经网络RNN实现:桨距控制控制风力发电机组研究