PyTorch中的原地操作

Question

我在想如何处理PyTorch中的原地操作问题。我记得在使用自动求导的情况下，原地操作经常会出问题。

实际上，我很惊讶下面的代码可以运行。虽然我没有测试过，但我相信在版本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 的积极缓存释放和重用使其非常高效，并且实际上很少有需要通过原地操作来降低内存使用的情况。除非你正在承受巨大的内存压力，否则可能永远不需要使用它们。

但即使它可以工作，在大多数情况下建议不要使用原地操作。

---

所以我的问题是：

• 使用原地操作会对性能产生多大影响？

• 在想要将张量的一个元素设置为特定值的情况下，如何避免使用原地操作？

谢谢！

- MBT

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

- Umang Gupta

Answer 2

这可能不是您问题的直接答案，而只是提供信息。

就计算图中的张量而言，原地操作可用于非叶张量。

叶张量是计算图的“末端”张量。根据官方文档（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.])

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

导致错误的示例：

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”覆盖）。我在这里找到了一个很好的图。

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

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

这仅仅是我的个人意见，请参考官方文档。

注意：

虽然原地操作对中间张量有效，但在进行一些原地操作时，尽可能使用 clone 和 detach 来显式地创建一个与计算图无关的新张量，以确保安全。

- Toru Kikuchi

Answer 3

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

- BBloggsbott