pytorch 指定 GPU

python pytorch 计算机视觉

看后台显卡使用情况(每 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]==，如下图所示

本文由 Yonghui Wang 创作，采用 知识共享署名4.0 国际许可协议进行许可
本站文章除注明转载/出处外，均为本站原创或翻译，转载前请务必署名
最后编辑时间为: Dec 19, 2024 12:13 pm