pytorch的自定义拓展之(二)——torch.autograd.Function完成自定义层

前言:前面的一篇文章中,已经很详细的说清楚了nn.Module、nn.functional、autograd.Function三者之间的联系和区别,虽然autograd.Function本质上是自定义函数的,但是由于神经网络、层、激活函数、损失函数本质上都是函数或者是多个函数的组合,所以使用autograd.Function依然可以达到定义层、激活函数、损失函数、甚至模型的目的,就像我们使用nn.Module是一样的,只不过更偏底层,稍微复杂一些而已,因为需要自己定义求导函数。

但是需要特别注意的是,对于复杂的层或者是网络,使用autograd.Function几乎是不可行的,因为我们需要重新定义反向求导规则即backward函数,而复杂层或者网络没办法写出每一个参数的导函数,或者是即便写出来也是异常复杂(因为链式求导法则再加上一些非线性函数的关系)所以一般不推荐使用autograd.Function去定义层,更不要去定义模型,但是一般定义一个较简单的函数还是可以的。

所以本文依然只会涉及到简单的定义操作,旨在帮助更好地理解autograd.Function的工作过程。关于autograd.Function的详细定义过程,我们可以参考前一篇文章:

pytorch的自定义拓展之(一)——torch.nn.Module和torch.autograd.Function

一、autograd.Function的实例

前面的文章都是使用的实例方法来重写forward和backward方法,下面来看一下如果使用静态方法,像Function类定义的那样,怎么实现。

1.1 重写Function类的静态方法

import torch
from torch.autograd import Function

# 类需要继承Function类,此处forward和backward都是静态方法
class MultiplyAdd(Function):  
                                                             
    @staticmethod                                  
    def forward(ctx, w, x, b):                 
        ctx.save_for_backward(w,x)    #保存参数,这跟前一篇的self.save_for_backward()是一样的
        output = w * x + b
        return output                        
         
    @staticmethod                                 
    def backward(ctx, grad_output):    #获取保存的参数,这跟前一篇的self.saved_variables()是一样的
        w,x = ctx.saved_variables  
        print("=======================================")             
        grad_w = grad_output * x
        grad_x = grad_output * w
        grad_b = grad_output * 1
        return grad_w, grad_x, grad_b  # backward输入参数和forward输出参数必须一一对应
 
x = torch.ones(1,requires_grad=True)  # x 是1,所以grad_w=1
w = torch.rand(1,requires_grad=True)  # w 是随机的,所以grad_x=随机的一个数
b = torch.rand(1,requires_grad=True)  # grad_b 恒等于1

print('开始前向传播')
z=MultiplyAdd.apply(w, x, b)   # forward,这里的前向传播是不一样的,这里没有使用函数去包装自定义的类,而是直接使用apply方法
print('开始反向传播')
z.backward()                   # backward

print(x.grad, w.grad, b.grad)
'''运行结果为:
开始前向传播
开始反向传播
=======================================
tensor([0.1784]) tensor([1.]) tensor([1.])
'''

注意:上面最大的不同除了使用的是静态方法以外,最大的不同在于,我没有使用一个函数去包装我的自定义类,而是直接使用了  z=MultiplyAdd.apply(w, x, b)  去完成前向运算过程,

这个apply方法是定义在Function类的父类_FunctionBase中定义的一个方法,但是这个方法到底是怎么实现的还不得而知。

这里到底是为什么,我还没有搞得特别清楚,因为我没有找到关于apply的详细代码所在,如果有大佬知道,望告知,万分感谢!

二、autograd.Function结合nn.Module

前面已经多次强调,虽然Function类可以用来定义模型,但是不要这么去做,因为Function类本身是为自定义函数而存在,我们在这里演示一下如何使用Function类自定义一个层,然后使用这个自定义的层来搭建网络。

为了简单,本文以搭建一个线性层作为实例说明:

参考下面的文章:

定义torch.autograd.Function的子类,自己定义某些操作,且定义反向求导函数

 

你可能感兴趣的:(深度学习,pytorch)