【亲测有用】如何使用pytorch指定GPU训练

实现方法

  • 1. 直接在程序中设置
  • 2. 在运行配置中设置(推荐)


摘要:深度学习中,跑模型代码的过程中,可能主GPU被占用而无法进行训练,这个时候,可知通过torch库来指定训练的GPU号。

1. 直接在程序中设置

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3':该命令是指定哪几块GPU可被系统识别,=前面是环境变量,表示CUDA中可被看到的设备,=后面是GPU序号,‘0,1,2,3’分别表示0、1、2、3号GPU。如果你只想用1、2、3号GPU,可将代码改成:

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'  # 这行代码须放在程序访问GPU之前,如程序的开头。

2. 在运行配置中设置(推荐)

(1)在pycharm中点击上方菜单栏:运行/编辑配置,如下所示:
【亲测有用】如何使用pytorch指定GPU训练_第1张图片
(2)跳出如下界面,然后在环境变量一栏添加上:CUDA_VISIBLE_DEVICES=1
【亲测有用】如何使用pytorch指定GPU训练_第2张图片
具体添加方法如下:
【亲测有用】如何使用pytorch指定GPU训练_第3张图片

(3)保存配置
【亲测有用】如何使用pytorch指定GPU训练_第4张图片
到这里便完成了。
现在,运行或调试都是在1号GPU上进行了

你可能感兴趣的:(pytorch,人工智能,python)