【Pytorch】基于torch.autograd.Function类自定义网络层

torch.autograd.Function是pytorch中用于自定义张量操作函数的类，通过实现forward和backward方法分别定义前向计算和反向求导过程，从而实现自定义网络层在整个动态流图中计算的传递。

先贴一张官方说明：

在自定义时，主要遵循如下流程和原则：
（1）forward和backward均为静态方法，其第一个参数必须为ctx（官方示例，可自定义名称）上下文管理器变量。在forward时将需保存给backward的值通过进行保存（ctx.save_for_backward）和获取（ctx.saved_tensors），注意两者数量的一致性；
（2）forward方法中的return值为该层网络传给下一层网络的张量，其数量和shape与backward中接收的变量一致，从而保证链式传到的可行；
（3）backward方法中的return值为BP过程中传给上一层网络的导数值，其数量和shape与forward方法中传入的张量数量和shape保持一致，从而证链式传导的可行；
（4）由上可知，一个完整的Function子类，需要完整的定义如下过程：张量数据从上一层网络传入——>本层计算——>传入下一层网络——>下一层网络反向传播的输入值——>本层求导过程——>传给上一层网络的反向传播链式求导结果；
（5）Function子类仅仅定义了张量的FN和BP运算过程，为了能作为网络层接入整体模型中，一般还需要定义对应的nn.Module子类，并在其中的forward中通过调用上述定义的Function子类的apply方法。

下面为一个简单的示例：
自定义一个网络层计算${\mathbf{x}}_1^3+{\mathbf{x}}_2^2$，最终输出为张量各元素的和。显然，反向传导的梯度值应当为$(3{\mathbf{x}}_1^2,2{\mathbf{x}}_2)$。

import torch
from torch.autograd import Function
import torch.nn as nn

# 张量函数的前向和反向过程
class SquareFun(Function):
    
    @staticmethod
    def forward(ctx, x1, x2):
        result = x1 ** 3 + x2 ** 2
        ctx.save_for_backward(x1, x2)
        return result
    
    @staticmethod
    def backward(ctx, grad_out):
        x1, x2 = ctx.saved_tensors
        return 3 * x1**2 * grad_out, 2 * x2 * grad_out

# 定义对应的module层
class Square(nn.Module):
    def forward(self, x1, x2):
        return SquareFun.apply(x1, x2)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1 = Square()
    
    def forward(self, x1, x2):
        y = self.layer1(x1, x2)
        y = y.sum()
        return y

运行一个实例，验证了上述流程的正确性。

net = Net()
x1 = torch.tensor([1,2,3], requires_grad=True, dtype=torch.float32)
x2 = torch.tensor([4,5,6], requires_grad=True, dtype=torch.float32)
out = net(x1, x2)

out  # tensor(113., grad_fn=<SumBackward0>)
out.backward()

x1.grad  # tensor([ 3., 12., 27.])
x2.grad  # tensor([ 8., 10., 12.])