How can a customized two-tensor `autograd.Function` be backward twice w/o. `retain_graph=True`?

username: DicardoX
cooked: <p>Hi, I have a customized <code>torch.autograd.Function</code> like this:</p>
<pre data-code-wrap="python"><code class="lang-python">@staticmethod
def forward(ctx, input, weight, bias):
    # Input is a tuple with two tensors, each for an independent task
    # Concatenate along batch dimension
    split_r = [_t.size(dim=0) for _t in input]
    concat_tensor = torch.concat(tensors=[_t for _t in input], dim=0)

    # Fused forward
    concat_output = concat_tensor.mm(weight.t())
    if bias is not None:
        concat_output += bias.unsqueeze(0).expand_as(concat_output)

    # Split along batch dimension
    output = tuple(torch.split(concat_output, split_r, dim=0))
    # Save
    ctx.save_for_backward(*input, weight, bias)
    return output

@staticmethod
def backward(ctx, grad_output):
    saved_tensors = ctx.saved_tensors
    inputs = saved_tensors[:-2]
    (weight, bias) = saved_tensors[-2:]

    grad_input, grad_weight, grad_bias = None, None, None
    # Naive gradient computation
    if ctx.needs_input_grad[0]:
        grad_input = grad_output.mm(weight)
    if ctx.needs_input_grad[1]:
        # Not compute weight gradient
        grad_weight = None
    if bias is not None and ctx.needs_input_grad[2]:
        grad_bias = grad_output.sum(0)

    return grad_input, grad_weight, grad_bias
</code></pre>
<p>Then, I use the following code to test it:</p>
<pre data-code-wrap="python"><code class="lang-python">x_1 = torch.randn(size=(2, 2), dtype=torch.float32, device="cuda:0", requires_grad=True)
x_2 = torch.randn(size=(2, 2), dtype=torch.float32, device="cuda:0", requires_grad=True)
labels = [
    torch.randn(size=(2, 2), dtype=torch.float32, device="cuda:0", requires_grad=True) for _ in range(2)
]

# Shared layer
fusion_out = CustomLinear.apply((x_1, x_2))
o_1, o_2 = fusion_out

# Task-specific layer
layer_1, layer_2 = nn.Linear(...), nn.Linear(...)
y_1 = layer_1(o_1)
y_2 = layer_2(o_2)

# Loss
loss_1 = loss_fn(y_1, labels[0])
loss_2 = loss_fn(y_2, labels[1])

# Backward
torch.autograd.backward(loss_1, None)
torch.autograd.backward(loss_2, None)
</code></pre>
<p>However, when I backward the second loss, I got: <code>Error: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.</code></p>
<p>There are some other task-specific layers (e.g., <code>layer_1</code>/<code>layer_2</code>, much more than the shared layer) in my network (for each independent task, this <code>CustomLinear</code> is the ONLY part that these two tasks share) and I don’t want to (1) set all these layers <code>retain_graph=True</code> or (2) backward these two losses concurrently like <code>torch.autograd.backward(loss_1 + loss_2, None)</code>, how can I do? I’d appreciate any hint on how to do this.</p>

username: KFrank
cooked: <p>Hi Dicardo!</p>
<aside class="quote no-group" data-username="DicardoX" data-post="1" data-topic="205730" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dicardox/48/71184_2.png" class="avatar"> DicardoX:</div>
<blockquote>
<p>Hi, I have a customized <code>torch.autograd.Function</code></p>
</blockquote>
</aside>
<p>I don’t believe that the use of a custom <code>Function</code> is relevant to the<br>
question you are asking.</p>
<p>I believe you want to know how to backpropagate two task-specific loss<br>
values through two task-specific layers and then through a single shared<br>
layer without backpropagating twice or using <code>retain_graph = True</code>.</p>
<p>You would use the same method for doing this whether or not your layers<br>
include custom <code>Function</code>s.</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">x_1 = torch.randn(size=(2, 2), dtype=torch.float32, device="cuda:0", requires_grad=True)
x_2 = torch.randn(size=(2, 2), dtype=torch.float32, device="cuda:0", requires_grad=True)
...
# Backward
torch.autograd.backward(loss_1, None)
torch.autograd.backward(loss_2, None)
</code></pre>
<p>However, when I backward the second loss, I got: <code>Error: Trying to backward through the graph a second time</code><br>
…<br>
I don’t want to (1) set all these layers <code>retain_graph=True</code> or (2) backward these two losses concurrently like <code>torch.autograd.backward(loss_1 + loss_2, None)</code></p>
</blockquote>
</aside>
<p>I’m not sure what you mean by “concurrently” – at some point you do have<br>
to somehow backpropagate both losses.</p>
<p>The straightforward approach is to use the feature where you can pass<br>
multiple losses (not summed together) to <code>autograd.backward()</code>.</p>
<p>Here is a script that illustrates this:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2024)

x_1 = torch.randn (2, 2, requires_grad = True)
x_2 = torch.randn (2, 2, requires_grad = True)

# shared layer
t = x_1 @ x_2
o_1 = t @ x_1
o_2 = x_2 @ t

# task-specific layers
layer_1 = torch.nn.Linear (2, 3)
layer_2 = torch.nn.Linear (2, 3)
y_1 = layer_1 (o_1)
y_2 = layer_2 (o_2)

# losses
loss_1 = y_1.sum()
loss_2 = y_2.sum()

# use retain_graph to compute grads and save
loss_1.backward (retain_graph = True)
loss_2.backward()
x_1_grad = x_1.grad.clone()
x_2_grad = x_2.grad.clone()

# clear grads
x_1.grad = None
x_2.grad = None
layer_1.zero_grad()
layer_2.zero_grad()

# repeat forward pass
t = x_1 @ x_2
o_1 = t @ x_1
o_2 = x_2 @ t
y_1 = layer_1 (o_1)
y_2 = layer_2 (o_2)
loss_1 = y_1.sum()
loss_2 = y_2.sum()

# use one autograd.backward() to backpropagate both loss_1 and loss_2
torch.autograd.backward ((loss_1, loss_2))

print ('torch.allclose (x_1.grad, x_1_grad):', torch.allclose (x_1.grad, x_1_grad))
print ('torch.allclose (x_2.grad, x_2_grad):', torch.allclose (x_2.grad, x_2_grad))
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.3.1
torch.allclose (x_1.grad, x_1_grad): True
torch.allclose (x_2.grad, x_2_grad): True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: DicardoX
cooked: <p>Sincere thanks to your detailed answer, Frank!</p>
<p>I got your solution by passing multiple losses to <code>autograd.backward()</code> at once and I believe it will work. Actually, the case <code>backward(loss_1 + loss_2, None)</code> is just a naive illustration and my current solution to this is indeed something like you suggest:</p>
<pre data-code-wrap="python"><code class="lang-python"># First, backward each loss seperately to get the input gradient to the output layer (i.e., the last layer before loss computation)
torch.autograd.backward(loss_1, None)
torch.autograd.backward(loss_2, None)

# As in forward pass, I have detached the input tensor to the last layer by: 
# input_1_detached = input_1.detach()
# input_2_detached = input_2.detach()
# loss_1 = loss_fn(output_layer(input_1_detached), labels[0])
# loss_2 = loss_fn(output_layer(input_2_detached), labels[1])
# 
# After I backward the losses, the gradients will be accumulated on input_1_detached/input_2_detached

# Then, I continue to backward the rest of the network by:
# Concatenate the gradients and the input tensors
concat_tensor = torch.concat([input_1, input_2])
concat_grad = torch.concat([input_1_detached.grad, input_2_detached.grad])
# Backward pass
torch.autograd.backward(concat_tensor, grad_tensors=concat_grad)
</code></pre>
<p>In this case, I can backward two losses w/o. any accuracy loss. Unfortunately, currently in my task schedule, I must enable detachable backward passes of each task (i.e., first backward loss_1, then backward loss_2). Therefore, the implementation above hinders the flexible schedule of the backward computation.  Maybe I need to implement a customized autograd engine as an extension to PyTorch.</p>
<p>Best.</p>
<p>Dicardo</p>

username: KFrank
cooked: <p>Hi Dicardo!</p>
<aside class="quote no-group" data-username="DicardoX" data-post="3" data-topic="205730" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dicardox/48/71184_2.png" class="avatar"> DicardoX:</div>
<blockquote>
<p>I must enable detachable backward passes of each task (i.e., first backward loss_1, then backward loss_2).</p>
</blockquote>
</aside>
<p>Yes, if your use case prohibits you from backpropagating <code>loss_1</code> and<br>
<code>loss_2</code> with a single <code>backward()</code>, your solution, where you use <code>.detach()</code><br>
to break the computation graph and then chain together <code>backward()</code> calls<br>
to piece the graph back together, is the appropriate method.</p>
<aside class="quote no-group">
<blockquote>
<p>Therefore, the implementation above hinders the flexible schedule of the backward computation.</p>
</blockquote>
</aside>
<p>Without knowing why your use case has the constraints you hint at, it’s<br>
not really possible to say more.</p>
<p>Best.</p>
<p>K. Frank</p>