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

发布时间:2024年01月19日

在 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 相关设置

????????这将在代码中明确指定程序使用的 GPU 设备。根据你的需求,你可以在 set_device 中传递不同的 GPU 索引。

文章来源:https://blog.csdn.net/weixin_51659315/article/details/135549448
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。