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

Question

问题总结：自定义函数的backward pass中如何处理输入和输出的维度？

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

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 。 它甚至打算考虑使用自定义函数处理批处理吗？

手册中似乎没有解决这些问题，并且提供的示例非常简单，根本没有帮助。 也许隐藏在某处的公式更详细地解释了提供的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的雅可比行列式dx_i/dy_j ）。

换句话说，雅可比行列式只是隐式使用，大部分不需要，因此为了记忆而取消。

对于批处理部分，请注意小批量学习大多只是平均梯度。 PyTorch 期望您在向后函数中处理它（同样，该函数尽可能少地返回并节省尽可能多的内存）。

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

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

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