PyTorch autograd：自定义函数梯度的维度？

Question

问题总结：自定义函数的反向传递中输入输出的维数是如何处理的？

根据manual，自定义函数的基本结构如下：

class MyFunc(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input): # f(x) = e^x
        result = input.exp()
        ctx.save_for_backward(result)
        return result

    @staticmethod
    def backward(ctx, grad_output): # df(x) = e^x
        result, = ctx.saved_tensors
        return grad_output * result

对于单个输入和输出维度，这非常好并且非常有效。但是对于更高的维度，向后传递变得混乱。显然，PyTorch 只接受 backward 的结果，它与 forward 的结果具有相同的维度（对于相同的输入）。返回错误的形状会产生 RuntimeError: Function MyFunc returned an invalid gradient at index 0 - got [*] but expected shape compatible with [*]。所以我想知道：向后实际计算的是什么？

它不是雅可比行列式？ 例如，当我有一个函数 f(x) = ( f_1(x_1, ... , x_n), ... , f_k(x_1, ... , x_n) ) 具有 n 输入和 k 输出时，我会预计梯度计算将产生维数 k*n 的雅可比矩阵。然而，PyTorch 实现只需要一个维度 n 的向量。那么向后的结果到底是什么意思，不可能是雅可比行列式？

而且它不处理批次？此外，如果我想通过这个函数推送一批输入向量怎么办，例如批量大小 b 的维度 b*n 的输入。然后，渐变也应该具有 b*n 的形状，而不是像 b*k*n 这样的东西。是不是还要考虑用自定义函数处理批次？

None 这些问题似乎在手册中得到了解决，提供的 examples 非常简单，根本没有帮助。也许某处隐藏着公式，更详细地解释了提供的Function界面的背景，但我还没有找到它们。

Answer 1

它不是 store/return 雅可比行列式（我想这与内存考虑有关）。

从训练的角度来看，我们不需要雅可比矩阵来进一步更新 parameters/back-propagating。

为了更新参数，我们只需要dL/dy_j、j<n:

y_j -= alpha * dL/dy_j

并且对于 z 的反向传播，说 z=f(y)=f(g(x)):

dL/dz_k = dL/dy_j * dy_j/dz_k

有人可能会说“但是我们这里需要 dy_j/dz_k！” -- 确实如此，但我们不需要存储它（就像我们在这一步中根本不使用 dx_i/dy_j 的雅可比行列式一样）。

换句话说，雅可比行列式仅被隐式使用，大部分情况下不需要，因此为了便于记忆而将其删除。

对于批处理部分，请注意小批量学习主要只是平均梯度。 PyTorch 希望您在后向函数中处理它（同样，这样函数 returns 尽可能少并尽可能节省内存）。

注意：可以“收集”雅可比矩阵并获得您提到的 n 大小的向量。具体来说，对 k 维度求和，然后对批处理维度取平均值。

编辑：不是 100% 确定，但我认为（f(x)=y）的后向调用预计 return 这个向量：

其中 \nabla x 是 backward 的输入参数。

PyTorch autograd：自定义函数梯度的维度？

PyTorch autograd: dimensionality of custom function gradients?

python

pytorch

autograd