1.使用方法
torch.autograd.grad如何使用
torch.autograd.grad(outputs,# 需要求导的目目标张量(如 Q 值),必须是标量或向量。inputs,# 需要对其求导的源张量(如动作 a),可以是任意形状grad_outputs=None,# 关键!当 outputs 是向量时,需要提供外部梯度(链式法则的初始值)retain_graph=None,# 是否保留计算图以供后续反向传播。默认 False,调用后释放图create_graph=False# 是否构建高阶导数图(用于二阶梯度)。默认 False)其中,如果是标量(如总损失),则grad_outputs可省略(默认为 1.0);如果是向量(如多个 Q 值组成的张量),则必须通过grad_outputs指定每个输出分量的权重系数,以完成向量-雅可比乘积(VJP)计算。
2. 通过例子理解
假如Y=f(X)Y = f(X)Y=f(X)
形状:Y[N,K]Y[N, K]Y[N,K]和X[N,M]X[N, M]X[N,M]
torch.autograd.grad(outputs=Y,inputs=X,grad_outputs=torch.ones_like(Y))[0]在使用torch.autograd.grad函数求梯度时,我的理解:
函数里填:
outputs=Y,
inputs=X,
grad_outputs=torch.ones_like(Y)
从函数外部的理解
其中grad_outputs=torch.ones_like(Y)代表对Y中每个元素求偏导后的加权。
最后得到的结果是:Y中每行,对X对应的行的元素,求偏导,再根据grad_outputs求加权和,得到的矩阵,形状是[N,M][N, M][N,M].
从函数内部的理解
函数里面的具体流程:
grad_outputs代表对Y中每个元素的加权。- 然后把Y每行的每个元素,对X对应的行的分量分别求偏导(Y和X全部展开,对应求导),得到一个形状是[N∗K,N∗M][N*K, N*M][N∗K,N∗M]的雅可比矩阵 J(块对角矩阵)。
grad_outputs=torch.ones_like(Y),和输出形状一样[N,K][N, K][N,K],经过内部展开成[N∗K,1][N*K, 1][N∗K,1],也就是vvv,代表对Y中所有元素对X求偏导之后的加权系数,实际上计算vT⋅Jv^{T} · JvT⋅J,得到的是形状为[1,N∗M][1, N*M][1,N∗M]的雅可比向量积(JVP),再由内部 reshape 变换成和输入XXX一样的形状[N,M][N, M][N,M]。
总结
经过这个函数,XXX和YYY的梯度并没有改变,而是输出了一个YYY对XXX求导的矩阵,是一个“无副作用的梯度查询器”,绝不污染任何变量的.grad属性。