如何在Python深度学习中设置GPU卡号

诸神缄默不语-个人CSDN博文目录

本文介绍在Python深度学习代码运行的过程中，如何设置GPU卡号（包括PyTorch和TensorFlow适用的写法），主要适用于单卡场景，以后可能会增加多卡场景。
常见适用场景：在多卡机子上，代码往往默认适用GPU 0，但有时需要使用1、2等其他GPU，因此需要手动设置。

1. CUDA_VISIBLE_DEVICES

设置之后，代码运行时就仅能看到这个被设置的GPU序号。如宏观逻辑号为1的GPU，设置后，代码运行时cuda:0就会直接将逻辑号为0的GPU定位到真实的1卡上。
就可以直接这么写：device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

1. 在代码中设置（注意需要写在深度学习代码之前）：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

2. 在代码运行命令之前设置：CUDA_VISIBLE_DEVICES=1 python run.py

2. PyTorch直接转移张量的device

一般来说输入就直接把每个张量都to(device)
模型中，已经注册好的张量，可以直接通过将模型实例to(device)就自动实现转换；而模型中未注册的张量（如在forward()等函数中新建的、辅助模型实现更多操作的张量）

3. 本文撰写过程中使用的参考资料

1. 使用CUDA_VISIBLE_DEVICES设置显卡_华科附小第一名的博客-CSDN博客_cuda_visible_devices