本文介绍在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")
- 在代码中设置(注意需要写在深度学习代码之前):
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
- 在代码运行命令之前设置:
CUDA_VISIBLE_DEVICES=1 python run.py
2. PyTorch直接转移张量的device
一般来说输入就直接把每个张量都to(device)
模型中,已经注册好的张量,可以直接通过将模型实例to(device)
就自动实现转换;而模型中未注册的张量(如在forward()
等函数中新建的、辅助模型实现更多操作的张量)