Pytorch指定GPU的方法总结

Pytorch指定GPU的方法

改变系统变量

1. 改变系统环境变量仅使目标显卡,编辑 .bashrc文件，添加系统变量

   export CUDA_VISIBLE_DEVICES=0
   #这里是要使用的GPU编号
   

2. 在程序开头设置

   os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
   

3. 在运行程序时指定

   # 运行程序时使用命令行,来设置该程序可见的gpu:
   CUDA_VISIBLE_DEVICES=0,1,2,3 python xxx.py
   

使用torch.cuda接口