Convert PyTorch tensor to CuPy array without detaching graph?

Ramansh_Sharma: <p>Hello,</p>
<p>I have a PyTorch tensor that is on a CUDA device and has a computation graph (it is built from another PyTorch CUDA tensor).</p>
<p>I want to convert this to a CuPy array without losing the computation graph. Is this possible at all?</p>
<p>When I try to convert the PyTorch tensor to CuPy array without using <code>detach</code> first, I get this error.<br>
<code>TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.</code></p>
<p>For reference, <a href="https://docs.cupy.dev/en/stable/user_guide/interoperability.html" rel="noopener nofollow ugc">these</a> are CuPy docs which recommend using DLPack to do this conversion (which loses the computation graph).</p>

Ramansh_Sharma: <p>Hello <a class="mention" href="/u/ptrblck">@ptrblck</a>. If possible, could you please help me with this? I would really appreciate the help. Thank you.</p>

ptrblck: <p>I think using <code>DLPack</code> sounds like the right approach, but I would also assume you would need to write a custom <code>autograd.Function</code> including the <code>backward</code> method as described <a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html">here</a> as Autograd would not be aware of CuPy operations.</p>

Ramansh_Sharma: <p>Thank you for the reply. I agree with the custom autograd function idea, we are indeed using something like that in our code. But since CuPy does not retain the graph of the PyTorch tensor, once we perform some operations in CuPy and convert back to PyTorch, we are unable to differentiate w.r.t. a PyTorch tensor that was used in creating PyTorch tensor that was converted to CuPy.</p>
<pre><code class="lang-python">x = torch.rand(10, 1).requires_grad_(True).to("cuda")
y = x ** 2

cupy_y = cupy.from_dlpack(y)

# perform some operations on cupy_y to get new_cupy_y

y_torch = torch.from_dlpack(new_cupy_y.toDlpack())

# this gives an error because y_torch is not made from x
torch.autograd.grad(y_torch, x, grad_outputs=torch.ones_like(y_torch), create_graph=True)
</code></pre>
<p>I think I will have to use PyTorch’s C++ autograd kernels in this case, where this problem shouldn’t be.</p>

ptrblck: <p>No, I don’t think you need to implement C++ Autograd kernels, but should be able to write a custom <code>autograd.Function</code> in Python defining the <code>forward</code> and <code>backward</code> methods as described in the link in my previous post.</p>

kmaehashi: <p>CuPy interoperability docs has a full working example of <code>torch.autograd.Function</code> implemented in CuPy <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>
<p><a href="https://docs.cupy.dev/en/latest/user_guide/interoperability.html#using-custom-kernels-in-pytorch" class="onebox" target="_blank" rel="noopener nofollow ugc">https://docs.cupy.dev/en/latest/user_guide/interoperability.html#using-custom-kernels-in-pytorch</a></p>
<p>This example uses raw CUDA kernel code, but the principle is the same when using CuPy’s builtin functions. Note that although <code>x.detach()</code> is called inside <code>forward()</code>, a computational graph will be constructed between inputs and outputs when running <code>y = CuPyLog.apply(x)</code>.</p>

Ramansh_Sharma: <p>Thank you for your replies. Giving this more thought, I have an idea to retry this by using the custom <code>autograd.Function</code> autograd kernel. Right now, I was unable to return the derivative w.r.t the CuPy matrix input to the <code>forward</code> method, but I can do it if I take as input a PyTorch tensor and use CuPy in my forward method regardless.</p>