python - PyTorch 自动毕业 : dimensionality of custom function gradients?

Question

问题总结:自定义函数的backward pass中如何处理输入和输出的维度？
根据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*n批量大小 b .然后，而不是像 b*k*n 这样的东西预计渐变也具有形状 b*n .它甚至打算考虑使用自定义函数处理批处理吗？
手册和提供的 examples 中似乎没有解决这些问题。非常简单，这根本没有帮助。也许隐藏在某处的公式可以解释所提供的背景 Function更详细的界面，但我还没有找到它们。

Answer 1

它不存储/返回雅可比行列式(我想这与内存考虑有关)。
从训练的角度来看，我们不需要雅可比矩阵来进一步更新参数/反向传播。
对于更新参数，我们只需要 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 期望您在向后函数中处理它(同样，该函数尽可能少地返回并节省尽可能多的内存)。
注意:可以“收集”雅可比矩阵并获得 n - 你提到的大小的向量。具体来说，对 k 求和批量维度上的维度和平均值。
编辑:不是 100% 确定，但我认为向后调用(f(x)=y)预计会返回这个向量:

哪里\nabla x是 backward 的输入参数.