指定GPU方法
方法一
在终端输入(以下均以指定使用gpu1为例)
export CUDA_VISIBLE_DEVICES=1
python train.py
或者直接
CUDA_VISIBLE_DEVICES=1 python train.py
方法二
在.py文件中加入如下代码
需要注意的是: 这一行代码要在import torch
前加入,不然可能无效
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
import torch
Bug!!
若其他文件也包含这行命令的话,import该文件时CUDA_VISIBLE_DEVICES
又会修改,所以需要注意
查看运行方法
终端输入
watch --color -n1 gpustat -cpu
需要提醒的是: 使用print(next(model.parameters()).device)
命令查看只会显示cuda:0表明是在gpu上运行,而不会显示哪一个gpu上运行。