I recently encountered a problem when I try to make a customized autograd Function on my own. In my case, it is quite different from conventional requirements, because the inputs to my Function need to be some lists/tuples of tensors.
To be more specific, my problem can be simplified as follows.
The input to my Function is a list of tensor (in the presented example, we may only consider one input list; but in fact, the Function may need to take more than one input lists).
x = [torch.randn(4,5), torch.randn(8,9)]
And let’s consider the simplest case, a customized ReLU Function.
I would suppose the following codes would run, but it won’t actually.
class MyReLU(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
#xneed to be a list of tensors
out = [torch.relu(t) for t in x]
ctx.save_for_backward(out)
return out
@staticmethod
def backward(ctx, y):
act, = ctx.saved_tensors
out = [(t1 > 0.0) * t2 for t1, t2 in zip(act, y)]
return out
I know the actual implementation might be different, but you may get the basic idea from the above codes. The above codes didn’t work, it seems that the Function only takes torch.Tensor as inputs.
I need to implement this because these tensors have distinct sizes and cannot be stacked together. Besides, this shown codes involves some for loop to iterate over all the tensors, but actually it won’t be necessary, since my real operation is written in CUDA C++ and this can be efficiently computed. And the length of the tensor list could be very large, iterating over all the tensors should be time-consuming and inefficient. So that’s why the input to the Function need to be a list of tensors.
Could you give any suggestions or workarounds about this? Thank you very much!!