PyTorch .to('cpu') 或 .to('cuda') 的文档

Question

我已经搜索了PyTorch文档，但找不到任何将张量移动到CPU或CUDA内存.to()

。我记得在某处看到，在nn.Module上调用to()是一种就地操作，但在张量上并非如此。

是否有适用于Tensor的就地版本？

我在哪里可以找到nn.Module和Tensor(可能还有其他地方)的to()文档？

Answer 1

您已经找到了文档！ 太好了。

.to不是张量的就地操作。但是，如果不需要移动，则返回相同的张量。

In [10]: a = torch.rand(10)
In [11]: b = a.to(torch.device("cuda"))
In [12]: b is a
Out[12]: False
In [18]: c = b.to(torch.device("cuda"))
In [19]: c is b
Out[19]: True

由于b已经在 GPU 上，因此不会进行任何更改c is b导致True.

但是，对于模型，它是一个就地操作，也会返回模型。

In [8]: import torch
In [9]: model = torch.nn.Sequential (torch.nn.Linear(10,10))
In [10]: model_new = model.to(torch.device("cuda"))
In [11]: model_new is model
Out[11]: True

将其保留在模型的位置是有意义的，因为模型的参数需要移动到另一个设备而不是model对象。对于张量，似乎创建了新对象。

Answer 2

我仍在学习浏览文档站点(搜索不是很好)。

这是我发现的：

torch.Tensor.to()

张量版本返回：

• self张量是否已经处于目标格式，或者，
• 目标格式的张量副本

未列出就地版本。

---

torch.nn.Module.to()

该文档暗示这是一个就地操作：

移动和/或强制转换参数和缓冲区。