指定GPU失败

几种指定GPU的方法

1、在命令行中先指定

export CUDA_VISIBLE_DEVICES=0,1,2,3

2、在主程序代码中使用os指定（在import torch之前）

import os 
os.environ['CUDA_VISIBLE_DEVICES'] = '2'

3、在命令行中指定

CUDA_VISIBLE_DEVICES=0,1,2,3 python xxx.py

*下面是没成功过的

4、在主程序代码中指定device时，同时指定gpu设备号

device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu' )

5、在主程序代码中使用torch指定

import torch
torch.cuda.set_device(0)

6、在主程序代码中使用torch指定

if __name__ == '__main__':
  with torch.cuda.device(2):
    main()

P.S.一些查看当前设备信息的torch代码

1. cuda是否可用
torch.cuda.is_available()
2. 返回gpu数量
torch.cuda.device_count()
3. 返回gpu名字
torch.cuda.get_device_name(0)
4. 返回当前设备索引
torch.cuda.current_device()