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

Question

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

我记得在某处看到调用to()ann.Module是就地操作，但在张量上却不是。

s有就地版本Tensor吗？

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

标签： pythonpytorch

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对象。对于张量，似乎创建了新对象。