Зарегистрируйте перехватчик для промежуточных узлов в PyTorch ⇐ Python

[Anonymous]

Представьте, что у вас есть что-то простое:

Код: Выделить всё

x = torch.tensor([4.0], requires_grad=True)
y = torch.tensor([2.0], requires_grad=True)
output = x * y + x / y

grad_x = torch.ones_like(output)
autograd.grad(output, x, grad_outputs=grad_x)

в результате чего получается такой вычислительный график:
[img]https://i. sstatic.net/9fEqzyKN.png[/img]

Теперь я хотел бы получить доступ к градиентам промежуточных узлов, то есть MulBackward0, DivBackward0 и AddBackward0. Я знаю, что PyTorch не сохраняет его по умолчанию. Я знаю, что могу сделать это явным образом, а затем определить save_grad и т. д.
Но есть ли возможность просто прикрепить перехватчик к любому из этих узлов без необходимости циклически проходить по графу? и т. д.?

Подробнее здесь: https://stackoverflow.com/questions/789 ... in-pytorch