Pytorch自定义算子反向传播

有关 自定义算子的实现前面已经提到，可以参考。本文讲述自定义算子如何前向推理+反向传播进行模型训练。

自定义一个线性函数算子

线性函数 $Y=X{W}^T+B$ 定义输入M 个X变量，输出N个Y变量的线性方程组。
$X$ 为一个 1 x M 矩阵，$W$为 N x M 矩阵, $B$ 为 1xN 矩阵，根据公式，输出$Y$为1xN 矩阵。其中 W 和 B 为算子权重参数，保存在模型中。
在训练时刻，模型输入$X$ , 和监督值$Y$，根据 算子forward（）计算的$Y^p$ ，计算Loss = criterion($Y$, $Y^p$ )，然后根据backward（）链式求导反向传播计算梯度值。最后根据梯度更新W 和 B 参数。

class LinearF(torch.autograd.Function):
    @staticmethod
    def symbolic(g, input, weight, bias):
        return g.op("MYLINEAR", input, weight, bias)
    @staticmethod
    def forward(ctx, input:Tensor, weight: Tensor, bias: Tensor) -> Tensor:
        output = input @ weight.T + bias[None, ...]
        ctx.save_for_backward(input, weight)
        return output
    @staticmethod
    def backward(ctx, grad_output:Tensor)->Tuple[Tensor, Tensor, Tensor]:
        # grad_output -- [B, N] = d(Loss) / d(Y)
        input, weight = ctx.saved_tensors
        grad_input = grad_output @ weight
        grad_weight = grad_output.T @ input
        grad_bias = grad_output.sum(0)
        # print("grad_input: ", grad_input)
        # print("grad_weight: ", grad_weight)
        # print("grad_bias: ", grad_bias)

        return grad_input, grad_weight, grad_bias

如何实现反向传播

前向推理比较简单，就根据公式来既可以。反向传播backward() 怎么写呢？
反向传播有两个输入参数，第一个为ctx，第二个grad_output，grad_output就是对forward() 输出output 的求导，如果是最后的节点，那就是loss对输出的求导，否则就是下一层对输出求导，输出grad_input, grad_weight, grad_bias则分别对应了forward的输入input、weight、bias的梯度。这很容易理解，因为是在做链式求导，LinearFunction是这条链上的某个节点，输入输出的数量和含义刚好相反且对应。
根据公式：
$Y=X{W}^T+B$
Loss = criterion($Y_{}^t$, $Y_{}$ )， 假设我们选择判别函数为L2范数，Loss = ∑ j = 0 N