# Retaining a subgraph of the original DAG

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.
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`.

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
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 !
``````

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