如何默认在GPU上运行PyTorch？

Question

我想使用cuda运行PyTorch。我对所有张量设置了model.cuda()和torch.cuda.LongTensor()。

如果我已经使用了model.cuda()，那么我是否必须明确地使用.cuda创建张量？

是否有一种方法可以默认在GPU上运行所有计算？

- tian tong

Answer 1

我认为你无法默认指定使用CUDA张量。不过，你应该查看一下PyTorch官方示例。

在ImageNet的训练/测试脚本中，他们使用了一个名为DataParallel的模型包装器。这个包装器有两个好处：

• 它处理多个GPU上的数据并行性
• 它处理将CPU张量转换为CUDA张量

如你在L164所见，你不必手动将输入/目标转换为CUDA张量。

请注意，如果你有多个GPU并且想使用其中一个，可以使用CUDA_VISIBLE_DEVICES前缀启动任何Python/PyTorch脚本。例如：CUDA_VISIBLE_DEVICES=0 python main.py。

- Remi

---

他们在第294行和296行。 - Muhammad Naufil

Answer 2

是的，你可以通过以下方式将默认的张量类型设置为cuda：

torch.set_default_tensor_type('torch.cuda.FloatTensor')

如果我已经使用了 model.cuda()，那么创建张量时是否需要显式使用 .cuda() 函数？

是的，不仅需要将模型的[参数]张量设置为 cuda，还需要将数据特征、目标（以及模型使用的任何其他张量）的张量设置为 cuda。

- iacob

Answer 3

您可以使用set_default_device。

对于Nvidia GPU：

torch.set_default_device('cuda')

如果您有多个GPU，您可以选择特定的一个：

torch.set_default_device('cuda:1')

对于CPU使用，请使用'cpu'

torch.set_default_device('cpu')

你可以通过创建一个简单的张量并获取其设备类型来检查默认设备: torch.tensor([1.2, 3.4]).device 如果选择了第一个 GPU，它会返回 device(type='cuda', index=0)。

---

你也可以使用设备来包装你的代码：

with torch.device('cuda:0'):
    t = torch.tensor([1.2, 3.4])

在这种情况下，默认值仅会针对包装代码进行更改。

- Simon

---

请注意，torch.set_default_device()仅适用于最新的pytorch版本（2.0+），在旧版本上无法使用。 - undefined

@KenvixZure 这是不是意味着它也适用于2.0版本，还是只适用于大于2.0的版本？ - undefined