跳到主要内容

PyTorch 基础用法

运算的内存开销

索引操作是不会开辟新内存的,而像 y = x + y 这样的运算是会新开内存的,然后将 y 指向新内存。

为了演示这一点,我们可以使用 Python 自带的 id 函数:如果两个实例的 ID 一致,那么它们所对应的内存地址相同;反之则不同。

import torch

x = torch.tensor([1, 2])
y = torch.tensor([3, 4])
id_before = id(y)
y = y + x
print(id(y) == id_before) # False

False

如果想指定结果到原来的 y 的内存,我们可以使用前面介绍的索引来进行替换操作。在下面的例子中,我们把 x + y 的结果通过 [:] 写进 y 对应的内存中。

x = torch.tensor([1, 2])
y = torch.tensor([3, 4])
id_before = id(y)
y[:] = y + x
print(id(y) == id_before) # True

True

我们还可以使用运算符全名函数中的 out 参数或者自加运算符 += (也即 add_()) 达到上述效果,例如 torch.add(x, y, out=y)y += x(y.add_(x))

x = torch.tensor([1, 2])
y = torch.tensor([3, 4])
id_before = id(y)
torch.add(x, y, out=y) # y += x, y.add_(x)
print(id(y) == id_before) # True