如何将所有张量设置为cuda设备?

Question

项目需要在GPU上进行计算，但手动切换每个张量到(设备)的时间太长。

我使用了这个，但是张量仍然留在cpu上。PIC有问题

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

Answer 1

上一个答案的方法名有误，应该是:

#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，您就不会受到影响。

Answer 2

要将所有张量设置为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)

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

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

这会在第一个CUDA设备上创建一个张量'x'。