指定GPU运行pytorch模型的n种方法

指定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上运行。