ICode9

精准搜索请尝试: 精确搜索
首页 > 其他分享> 文章详细

Pytorch使用autograd.function自定义op

2021-08-17 01:32:49  阅读:381  来源: 互联网

标签:function Function 自定义 autograd input Variable output backward grad


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读取


使用时调用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)

标签:function,Function,自定义,autograd,input,Variable,output,backward,grad
来源: https://www.cnblogs.com/Valeyw/p/15150354.html

本站声明: 1. iCode9 技术分享网(下文简称本站)提供的所有内容,仅供技术学习、探讨和分享;
2. 关于本站的所有留言、评论、转载及引用,纯属内容发起人的个人观点,与本站观点和立场无关;
3. 关于本站的所有言论和文字,纯属内容发起人的个人观点,与本站观点和立场无关;
4. 本站文章均是网友提供,不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属;如您发现该文章侵犯了您的权益,可联系我们第一时间进行删除;
5. 本站为非盈利性的个人网站,所有内容不会用来进行牟利,也不会利用任何形式的广告来间接获益,纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

专注分享技术,共同学习,共同进步。侵权联系[81616952@qq.com]

Copyright (C)ICode9.com, All Rights Reserved.

ICode9版权所有