PyTorch: When using backward(), how can I retain only part of the graph?
I have a PyTorch computational graph, which consists of a sub-graph performing some calculation, and the result of this calculation (let's call it x) is then branched into two other sub-graphs. Each of these two sub-graphs yields some scalar results (lets call them y1 and y2). I want to do a backward pass for each of these two results (that is, I want to accumulate the gradients of the two sub-graphs. I do not want to perform the actual optimization step).

Now, since memory is an issue here, I would like to perform the operations in the following order:
First, calculate x. Then, calculate y1, and perform y1.backward() while (and this is the key point) retaining the graph that leads to x, but freeing the graph from x to y1. Then, calculate y2, and perform y2.backward().

In other words, to save memory without sacrificing too much speed, I want to keep x without needing to recalculate it, but I want to drop all calculations leading from x to y1 after I don't need them anymore.

The problem is that the argument retain_graph of the function backward() will retain the entire graph leading to y1, whereas I need to retain only the part of the graph leading to x.

Here is an example of what I would ideally want:



import torch

w = torch.tensor(1.0)
w.requires_grad_(True)

# sub-graph for calculating `x`
x = w+10

# sub-graph for calculating `y1`
x1 = x*x
y1 = x1*x1
y1.backward(retain_graph=x) # this would not work, since retain_graph is a boolean and can either retain the entire graph or free it.

# sub-graph for calculating `y2`
x2 = torch.sqrt(x)
y2 = x2/2
y2.backward()


How can this be done?

The argument retain_graph will retain the entire graph, not just a sub-graph. However, we can use garbage collection to free unneeded parts of the graph. By removing all references to the sub-graph from x to y1, this sub-graph will be freed:



import torch

w = torch.tensor(1.0)
w.requires_grad_(True)

# sub-graph for calculating `x`
x = w+10

# sub-graph for calculating `y1`
x1 = x*x
y1 = x1*x1
y1.backward(retain_graph=True) # all graph is retained

# remove unneeded parts of graph. Note that these parts will be freed from memory (even if they were on GPU), due to python's garbage collection 
y1 = None
x1 = None

# sub-graph for calculating `y2`
x2 = torch.sqrt(x)
y2 = x2/2
y2.backward()

