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()  # 返回當前設備索引