PyTorch如何指定GPU - 果核剥壳

很多网友想知道PyTorch如何指定GPU，在使用PyTorch进行深度学习任务时，合理地利用GPU资源可以显著提高训练速度和性能。PyTorch提供了简便的方法来指定特定的GPU设备用于模型的训练和推断。通过将计算任务分配给适当的GPU，我们可以加快计算速度并利用GPU的并行计算能力。今天机务诶大家分享两种Pytorch指定GPU的方法，希望能够帮助到有需要的网友。

Pytorch指定GPU
方法一：使用CUDA_VISIBLE_DEVICES环境变量
这种方法是在运行PyTorch程序之前，设置一个环境变量CUDA_VISIBLE_DEVICES，它的值是一个或多个GPU的编号，用逗号分隔。例如，如果你想要使用第0和第2个GPU，你可以在命令行中输入：

export CUDA_VISIBLE_DEVICES=0,2
然后运行你的PyTorch程序，它会自动使用这两个GPU，并且把它们映射为第0和第1个设备。你可以通过torch.cuda.device_count()函数来查看可用的GPU数量，以及通过torch.cuda.current_device()函数来查看当前的GPU编号。

方法二：使用torch.cuda.set_device()函数
这种方法是在PyTorch程序中，使用torch.cuda.set_device()函数来指定使用哪一个GPU。这个函数接受一个整数参数，表示GPU的编号。例如，如果你想要使用第1个GPU，你可以在程序中添加如下代码：

import torch
torch.cuda.set_device(1)
注意，这个函数必须在创建任何张量或模型之前调用，否则会报错。另外，这个函数只对当前进程有效，如果你使用多进程来并行运行多个模型，你需要在每个进程中都调用这个函数，并且指定不同的GPU编号。

总而言之，我们可以根据自己的需求和菜做习惯来选择其中一种方法来进行设置。