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!