PyTorch指定GPU号

Pytorch指定GPU号
在PyTorch中，你可以通过指定GPU的ID来在特定的GPU上运行你的模型。这是通过使用 torch.cuda.set_device() 函数来实现的，该函数会将之后需要运行在GPU上的Tensor数据转移到指定的GPU上。
首先，你需要知道你的系统中可用的GPUID。你可以使用以下的代码来列出可用的GPUID：

import torch
print(torch.cuda.device_count())

假设你有两个GPU，那么运行上述代码将输出 2。这表示你可以在0和1这两个GPU上运行你的PyTorch代码。
要指定GPUID，你可以使用 torch.cuda.set_device() 函数。例如，如果你想在第二个GPU（即GPUID为1）上运行你的代码，你可以使用以下代码：

torch.cuda.set_device(1)

现在，任何你创建的Tensor都将默认在指定的GPU上。例如：

x = torch.randn(1)
print(x.device)  # 输出：device('cuda:1')

你也可以在创建Tensor时指定其所在的设备。例如：

x = torch.randn(1, device=torch.device('cuda:1'))
print(x.device)  # 输出：device('cuda:1')

请注意，一旦你指定了特定的GPU，该设备上的所有Tensor都将保留在该设备上，即使你之后创建的Tensor没有明确指定设备。如果你想让Tensor在不同的设备之间移动，你需要明确指定设备。
同样需要注意的一点是，并非所有的操作都可以在所有的GPU上进行。一些操作可能会依赖于特定的计算能力或硬件特性，因此你可能需要根据你的具体需求和硬件来选择合适的GPU。