一、准备工作
查看GPU状态和信息,找到空闲的GPU:
nvidia-smi
二、指定单GPU
从图中 Processes 表格我们可以发现0、1、2号GPU均是可以利用的。于是我们可以在python文件中加入以下代码,表示使用0号GPU:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0" # 使用0号GPU
或者可以在命令行窗口处输入,表示使用1号GPU运行代码:
CUDA_VISIBLE_DEVICES=1 python your_model.py # 使用1号GPU运行代码
再或者使用 torch.cuda.set_device()函数
指定gpu使用编号:
import torch
torch.cuda.set_device(2) # 使用2号GPU
三、指定多GPU
若此时所有GPU都是空闲的,那么最简单的做法是:
import torch
model = torch.nn.DataParallel(model) # 使用默认所有存在的GPU
而假如像上文所述只有0、1、2号显卡空闲,那么可以使用以下代码指定多个GPU:
(注意:无论是指定单个GPU还是多个GPU,设定 os 的语句需要放在 import torch 语句之前,否则指定GPU的操作可能无法生效)
import os
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, [0,1,2])) # 一般在程序开头设置
import torch
model = torch.nn.DataParallel(model) # 等价于os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2'
可以看到两者相比,后者其实就是限定了GPU范围。
或者直接使用更简便的方法:文章来源:https://www.toymoban.com/news/detail-443530.html
import torch
torch.nn.DataParallel(model, device_ids=[0, 1, 2]) # 使用0、1、2号GPU
四、补充
为GPU设置随机种子文章来源地址https://www.toymoban.com/news/detail-443530.html
import torch
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
到了这里,关于在服务器上指定GPU跑代码的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!