PyTorch：当使用backward()时，我怎样才能只保留图形的一部分？答案

【问题标题】：PyTorch: When using backward(), how can I retain only part of the graph?PyTorch：当使用backward()时，我怎样才能只保留图形的一部分？
【发布时间】：2018-06-08 15:37:46
【问题描述】：

我有一个 PyTorch 计算图，它由一个执行某些计算的子图组成，这个计算的结果（我们称之为x）然后分支到另外两个子图。这两个子图中的每一个都会产生一些标量结果（我们称它们为y1 和y2）。我想对这两个结果中的每一个做一个反向传递（即我想累积两个子图的梯度。我不想执行实际的优化步骤）。

现在，由于内存是一个问题，我想按以下顺序执行操作：首先，计算x。然后，计算y1，并执行y1.backward() 同时（这是关键点）保留通向x 的图形，但将图形从x 释放到y1 .然后，计算y2，执行y2.backward()。

换句话说，为了在不牺牲太多速度的情况下节省内存，我想保留x而不需要重新计算它，但我想在我不这样做之后放弃从x到y1的所有计算不再需要它们了。

问题是函数backward() 的参数retain_graph 将保留通向y1 的整个图形，而我只需要保留通向x 的图形部分。

这是我理想中想要的示例：

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()

如何做到这一点？

【问题讨论】：

标签： pytorch

【解决方案1】：

参数retain_graph 将保留整个图，而不仅仅是一个子图。但是，我们可以使用垃圾收集来释放图形中不需要的部分。通过删除从x 到y1 对子图的所有引用，该子图将被释放：

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()

【讨论】：