По умолчанию градиенты сохраняются только для параметров, в основном только для экономии памяти. Если вам нужны градиенты внутренних узлов графа вычислений, вам нужно иметь соответствующий тензор перед вызовом backward()
и добавить хук, который будет выполняться при обратном проходе.
Минимальное решение от Форум PyTorch :
yGrad = torch.zeros(1,1)
def extract(xVar):
global yGrad
yGrad = xVar
xx = Variable(torch.randn(1,1), requires_grad=True)
yy = 3*xx
zz = yy**2
yy.register_hook(extract)
#### Run the backprop:
print (yGrad) # Shows 0.
zz.backward()
print (yGrad) # Show the correct dzdy
В этом случае градиенты сохраняются в глобальной переменной, где они сохраняются после того, как PyTorch избавляется от них в самом графике.