PyTorch无法指定GPU的问题解决

最近在一台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()  # 返回当前设备索引