如何将所有张量设置到cuda设备？

Question

项目需要在GPU上计算，但是手动切换每个tensor.to(device)时间太长

我用过这个，但张量仍然保留在 cpu 上。

if torch.cuda.is_available():
     torch.set_default_tensor_type(torch.cuda.FloatTensor)

Answer 1

要将所有张量设置为 CUDA 设备，可以使用“torch”张量库的“to”方法。 to 方法允许您指定要将张量“移动到”的设备。例如，要将所有张量移动到第一个 CUDA 设备，可以使用以下代码：

    import torch

# Set all tensors to the first CUDA device
device = torch.device("cuda:0")
torch.set_default_tensor_type(device)

或者，您也可以在使用“设备”参数创建新张量时指定设备。例如：

import torch

# Set all tensors to the first CUDA device
device = torch.device("cuda:0")
x = torch.zeros(10, device=device)

这将在第一个 CUDA 设备上创建一个张量“x”。

Answer 2

之前的回答方法名有误，应该是：

#instead of:
torch.set_default_tensor_type(device)
#should be:
torch.set_default_device(device)

我把它安排成一个很好的形式，以便回退到 cpu 或第一个 cuda 设备。

device = torch.device(
                      f'cuda:{torch.cuda.current_device()}')
                      if torch.cuda.is_available()
                      else 'cpu'

torch.set_default_device(self.device)

这样，如果机器没有可用的 GPU，您就可以得到保障。