在 PyTorch 中，怎么指定程序使用的 GPU。

最新推荐文章于 2024-07-08 11:10:59 发布

阿罗的小小仓库

最新推荐文章于 2024-07-08 11:10:59 发布

阅读量1.2k

点赞数 9

CC 4.0 BY-SA版权

文章标签： pytorch 人工智能 python

本文介绍了如何在PyTorch中通过CUDA_VISIBLE_DEVICES环境变量指定程序使用的GPU，包括单GPU和多GPU设置，并提醒了torch.cuda.is_available()的检查以及在代码中动态指定GPU的示例。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在 PyTorch 中，你可以使用 CUDA_VISIBLE_DEVICES 环境变量来指定程序使用的 GPU。在你的命令行中运行脚本之前，你可以设置该环境变量。以下是如何指定显卡的一种方法：

假设使用第一张 GPU，可以运行以下命令：

CUDA_VISIBLE_DEVICES=0 python main.py

这将使程序只能看到第一张 GPU，并使用它进行计算。如果想使用多个 GPU，你可以在等号后面提供逗号分隔的 GPU 索引列表：

CUDA_VISIBLE_DEVICES=0,1 python main.py

在这个例子中，程序将只能看到索引为 0 和 1 的两张 GPU，并使用它们进行计算。

记得在运行之前确保 PyTorch 可以使用 CUDA，即 torch.cuda.is_available() 返回 True。在你的代码中，你已经使用了 args.cuda = torch.cuda.is_available() 来设置这个变量，确保 CUDA 可用性。

请注意，使用这种方法只是在运行脚本时指定程序能够看到哪些 GPU，而不是在代码中直接指定使用某一张 GPU。如果想在代码中指定使用特定的 GPU，可以在代码中添加相应的 PyTorch CUDA 设置，例如：

# 在你的代码中添加以下部分
if args.cuda:
    torch.cuda.set_device(0)  # 设置使用第一张 GPU
    # 其他 CUDA 相关设置