Hello,
I’m trying to implement a new Pytorch autograd function which should use a custom-derived gradient and combine it with a gradient obtained by the autograd engine. In essence I would like to do the following:
class some_funny_function(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
output1 = some_custom_function(input)
input = torch.tensor(input, requires_grad=True)
output2 = some_pytorch_function(input)
ctx.save_for_backward(input, output1, output2)
return output1, output2
@staticmethod
def backward(ctx, grad_input1, grad_input2):
input, output1, output2 = ctx.saved_tensors()
grad_output1 = custom_gradient(output1, grad_input1)
output2.backward(grad_input2)
grad_output2 = input.grad
return grad_output1 + grad_output2
I have some input that is processed by two different functions, one I can directly write with PyTorch functions. The other can’t be done with PyTorch because it involves calling third-party software. Fortunately, I have an analytic expression for this gradient which I would like to seamlessly integrate into a PyTorch module.
Since the computation that I abbreviated ba some_pytorch_function is rather lengthy and the gradient complicated to derive, I would much rather use the autograd engine to do the hard work for me.
Unfortunately, ctx seems not to be able to do the trick for me since requires_grad seems to be disabled in a custom autograd.Function.
Do you have any suggestions on how to combine autograd and manual gradients in a nice way? Any help would be much appreciated!
Thanks
Nerrror