Retaining a subgraph of the original DAG

1

Is there a way to preserve a subgraph of the DAG which does not contains its root? As an example, consider the following code:

x1 = torch.randn(10, requires_grad=True)
x2 = x1 ** 2.
y1 = torch.randn(10, requires_grad=True)
y2 = y1 + 1.
z = (x2 + y1).sum()
z.backward(retain_graph=True)

I need to reuse x2, for other subsequent computation. To this end, I want to retain the subgraph associated to x2 and x1, zeroing its gradients but keeping their saved tensors for the intermediate computation, while I want to delete the subgraph associated to z, y2 and y1.

2

You can do exactly what you did here and reuse x2 in another computation.
A graph is kept only as long as something can reference it.

x1 = torch.randn(10, requires_grad=True)
x2 = x1 ** 2.
# Let call Graph 1 the part of the graph that give x2
y1 = torch.randn(10, requires_grad=True)
y2 = y1 + 1.
# Lets call Graph 2 the part of the graph that give y2
# Note that the only way to access this graph is through y2
z = (x2 + y2).sum()
# Let call Graph 3 the part of the graph that give z
# Note that Graph 3 contains Graph 1 and Graph 2 as subgraphs !
# The only way to access Graph 3 is via z
# Now the only ways to acces graph 2 is z and y2
z.backward(retain_graph=True)
# Here all graph still exists.
del z
# Now Graph 3 is gone as z is gone, but 1 and 2 are still there
del y 
# Now Graph 2 is gone as well as both z and y2 are gone and only Graph 1 remains
other_stuff(x2)
# You can still call backward on Graph 1 !
3

True, I was erroneously using the nvidia-smi command to check the memory and I wasn’t detecting any change. I then discovered the torch.cuda.memory_allocated() function. Thanks!

1 Like