在服务器上有多块gpu时,如果不指定就会默认用cuda0,如何指定用哪一块GPU。
1、直接在代码的开头指定
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
2、用.to()或者.cuda()指定
device = torch.device("cuda:6" if torch.cuda.is_available() else "cpu") # 单GPU或者CPU
model.to(device)
net = net.cuda(6)
# .to()和.cuda()的区别在于.to()可以指导用cpu还是gpu,而.cuda()只能指定gpu
PS 1和2不要同时用。之前写代码的时候因为一直有数据不在同一块gpu上的问题,就想着上双保险,1和2都用了,结果直接在.cuda(6)的地方报错没有gpu6。排查了好久才发现是因为用1指定了gpu6之后,gpu6就被认为是cuda0了,并且只指定了一块gpu就被认为仅有这块gpu,后面再写.cuda(6)的时候就被认为是只有cuda0 的情况下指定cuda6 ,于是就报错了。
PPS 查看gpu使用情况
1)英伟达自带的命令 nvidia-smi
2) 借助gpustat这个包
如果没有这个包,先用 pip install gpustat 命令安装
gpustat -c / -u