无需复制即可交换 PyTorch 张量的正确方法

Question

我有两个具有相同尺寸的 PyTorch 张量 x、y。我想交换两个张量后面的数据，理想情况下无需复制。这样做的目的是在其他地方拥有保存张量 x 的代码，以便现在读取和写入数据 y，反之亦然。

正确的做法是什么？这种交换是否会破坏任何保留对这些张量的引用的代码？交换后张量的引用计数/GC 是否仍能正常工作？

Answer 1

首先，您应该知道 PyTorch 中的每个张量都有一个底层存储来保存其实际数据。您可以使用

.storage()

检索张量的底层存储。然后，您应该使用

.set_()

将张量的底层存储替换为新存储。

x_storage = x.storage()
y_storage = y.storage()

x.set_(y_storage)
y.set_(x_storage)

注意：交换过程不会影响对张量本身的引用。此外，交换存储不会干扰引用计数/GC，因为 PyTorch 会自动处理引用计数和垃圾收集。另一个原因是您没有创建新的张量或直接修改引用计数。