How do I do custom gradient calculation for a PyTorch Tensor-like type?

user_3741677 0
Viewed 142

I am attempting to create a tensor-like class that I wish to integrate with PyTorch. The class will have to store the tensor values and gradients in a custom format. I've been following the instructions at extending torch with a Tensor-like type. I have questions especially pertaining to gradient storage and calculation:

  1. I want to initialize my class from a (float) tensor and be able to convert it back. I know I can retrieve the values from the tensor using the numpy() function, but how do I get gradient data if I wish to store that too? Is directly accessing tensor.grad the correct way to do so? When I convert back to tensor, how can I give it back the stored gradient data?
  2. In the above link I saw the instructions for making functions like torch.add work with my custom type, but that seems to only deal with the forward pass and I will also need to modify how the gradient is calculated by autograd. How do I define both the custom forward and backward versions of torch.add? Elsewhere, I saw the instructions for extending torch.autograd with a custom function but I am not sure if those instructions can be merged with this case. Is there a way I can create a custom add function and say that that function's apply implements torch.add for my class?

Thank you

tensor pytorch python
1 Answers
  1. Use tensor.grad.data (check if grad not None!)
  2. Can you perform the operation on the torch.Tensor directly or does it have to be your custom type? Then the tensor could carry the autograd information for most function (you would only compute the gradient yourself on custom functions). You could also consider subclassing torch.Tensor, that would alleviate a lot of those problems.
user_21695990
Related