在 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 索引。