[人工智能]PyTorch指定GPU

看后台显卡使用情况(每1s刷新一次) watch -n 1 nvidia-smi

1. 设置可见GPU

注意！！！！！！！！！！！！！！
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2"此条命令运行必须放在import torch之前，否则不能生效！！！！！！！！！！！！！！！！！！！此条坑了我一上午

当然还是建议这条方法或者第二条，这样设置了以后对程序可见的卡就固定在这几张了，不会对他人正在使用的卡产生任何影响

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2"

2. 终端运行时设定

CUDA_VISIBLE_DEVICES=1,2,3 python train.py

3. 代码中设定

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3,4,5"


#将模型放入GPU中
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model, device_ids=[0,1,2,3])

等价于

if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model, device_ids=[2,3,4,5])

这是因为在第一段中已经设置了可见GPU为2，3，4，5，就是说假设我们有八块GPU，序号分别为0,1,2,3,4,5,6,7，本来他们的编号也分别为0,1,2,3,4,5,6,7，但是因为设置了可见GPU为2,3,4,5所以我们选取的GPU的序号为2,3,4,5，但是对于程序来说他们的编号分别为[0,1,2,3]，如下图所示