最近在一台8卡服务器上跑PyTorch的训练代码,一直都没有问题。但是重新安装cuda之后,无法指定在哪几块GPU上运行了,只能按照顺序从第0块开始使用。查了一些资料,问题得到解决。

1、在python程序中指定在哪几块GPU上运行,通常采用如下方式:

import os
import torch

os.environ["CUDA_VISIBLE_DEVICES"] = "4,5,6,7"

或者直接在命令行中执行以下命令(不推荐):

export CUDA_VISIBLE_DEVICES=4,5,6,7

2、按照之前的写法,突然上面的代码无效了,无论怎么修改可见的GPU编号,最后程序都是按照顺序从第0块开始使用。问题出在指定GPU这行代码的位置上面,需要将“os.environ["CUDA_VISIBLE_DEVICES"] = "4,5,6,7"”移到import torch等代码之前,紧随import os之后,即按照如下的方式:

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "4,5,6,7"

import torch

3、附带查看GPU信息的一些常用指令以备后用,如下:

import torch

torch.cuda.is_available()  # 查看cuda是否可用

torch.cuda.device_count()  # 返回GPU数目

torch.cuda.get_device_name(0)  # 返回GPU名称,设备索引默认从0开始

torch.cuda.current_device()  # 返回当前设备索引

 

 

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐