使用 PyTorch 进行就地操作

Posted 2023-03-12

【中文标题】使用 PyTorch 进行就地操作

【英文标题】：In-place operations with PyTorch

【发布时间】：2019-01-19 22:30:39

【问题描述】：

我想知道如何在 PyTorch 中处理就地操作。我记得在 autograd 中使用就地操作一直存在问题。

实际上，我很惊讶下面的这段代码可以正常工作，即使我没有测试过，我相信这段代码会在版本0.3.1 中引发错误。

基本上我想做的是将张量向量的某个位置设置为某个值，如下所示：

my_tensor[i] = 42

工作示例代码：

# test parameter a
a = torch.rand((2), requires_grad=True)
print('a ', a)
b = torch.rand(2)

# calculation
c = a + b

# performing in-place operation
c[0] = 0
print('c ', c)
s = torch.sum(c)
print('s ', s)

# calling backward()
s.backward()

# optimizer step
optim = torch.optim.Adam(params=[a], lr=0.5)
optim.step()

# changed parameter a
print('changed a', a)

输出：

a  tensor([0.2441, 0.2589], requires_grad=True)
c  tensor([0.0000, 1.1511], grad_fn=<CopySlices>)
s  tensor(1.1511, grad_fn=<SumBackward0>)
changed a tensor([ 0.2441, -0.2411], requires_grad=True)

显然在0.4.1 版本中。这工作得很好，没有警告或错误。

参考文档中的这篇文章：autograd-mechanics

在 autograd 中支持就地操作是一件困难的事情，我们 在大多数情况下不鼓励使用它们。 Autograd 的激进缓冲区 释放和重用使其非常高效，并且很少 就地操作实际上会降低任何内存使用量的情况 数额可观。 除非您在大量内存下操作 压力，你可能永远不需要使用它们。

但即使它有效，在大多数情况下不鼓励使用就地操作。

---

所以我的问题是：

就地操作的使用对性能的影响有多大？

在我想将张量的一个元素设置为某个值的情况下，如何使用就地操作？

提前致谢！

【参考方案1】：

我不确定就地操作对性能的影响有多大，但我可以解决第二个查询。您可以使用掩码而不是就地操作。

a = torch.rand((2), requires_grad=True)
print('a ', a)
b = torch.rand(2)

# calculation
c = a + b

# performing in-place operation
mask = np.zeros(2)
mask[1] =1
mask = torch.tensor(mask)
c = c*mask
...

【参考方案2】：

这可能不是您问题的直接答案，仅供参考。

就地操作适用于计算图中的非叶张量。

叶张量是作为计算图“末端”的张量。官方（来自is_leaf属性here），

对于 requires_grad 为 True 的张量，如果它们是由用户创建的，它们将是叶张量。这意味着它们不是操作的结果，因此 grad_fn 为 None。

没有错误的示例：

a = torch.tensor([3.,2.,7.], requires_grad=True)
print(a)   # tensor([3., 2., 7.], requires_grad=True)
b = a**2
print(b)   # tensor([ 9.,  4., 49.], grad_fn=<PowBackward0>)
b[1] = 0
print(b)   # tensor([ 9.,  0., 49.], grad_fn=<CopySlices>)
c = torch.sum(2*b)
print(c)   # tensor(116., grad_fn=<SumBackward0>)
c.backward()
print(a.grad)  # tensor([12.,  0., 28.])

另一方面，就地操作不适用于 leaf 张量。

导致错误的示例：

a = torch.tensor([3.,2.,7.], requires_grad=True)
print(a) # tensor([3., 2., 7.], requires_grad=True)
a[1] = 0
print(a) # tensor([3., 0., 7.], grad_fn=<CopySlices>)
b = a**2
print(b) # tensor([ 9.,  0., 49.], grad_fn=<PowBackward0>)
c = torch.sum(2*b)
print(c) # tensor(116., grad_fn=<SumBackward0>)
c.backward()  # Error occurs at this line. 

# RuntimeError: leaf variable has been moved into the graph interior

我认为上面第一个示例中的b[1]=0 操作并不是真正的就地操作。我想它用“CopySlices”操作创建了一个新的张量。就地操作之前的“旧 b”可能会保留在内部（只是它的名称被“新 b”覆盖）。我发现了一个不错的图here。

旧 b ---(CopySlices)----> 新 b

另一方面，张量a 是一个叶张量。在 CopySlices 操作a[1]=0 之后，它变成了一个中间张量。为了避免反向传播时叶子张量和中间张量之间的复杂混合，禁止对叶子张量进行CopySlices操作与反向并存。

这只是我个人的看法，请以官方文档为准。

注意：

虽然就地操作适用于中间张量，但在执行一些就地操作时尽可能使用克隆和分离来显式创建独立于计算图的新张量是安全的。

【参考方案3】：

对于您的第二个查询，当您执行c[i] = i 或类似操作时，通常会调用__setitem__。要使该操作就地进行，您可以尝试调用__setitem__ 函数（如果这是执行c[i] = i 操作的原因。