pytorch 进行深度学习训练模型
服务器上有多个GPU,需要制定代码在其中某个GPU运行时,可以采用下面的方式进行指定
目录
查看指定GPU的id,见以下
1、查看服务器上的gpu信息
服务器上的终端上输入
watch -n 0.5 nvidia-smi
2、获取id,这里取的是id 1
详情见
查看显卡及GPU相关信息,指定GPU、CPU运行 - JY小脚丫 - CSDN博客
https://blog.csdn.net/xu380393916/article/details/90452912
3、在代码的开始位置添加以下的代码
import torch
torch.cuda.set_device(1) # 指定gpu1
参考文献
(原)PyTorch中使用指定的GPU - darkknightzh - 博客园
https://www.cnblogs.com/darkknightzh/p/6836568.html