位置: IT常识 - 正文
1. 使用CUDA_VISIBLE_DEVICES
2. 使用cuda()和torch.cuda.set_device()
3. 使用device
4. 使用torch.nn.DataParallel
推荐整理分享在GPU上运行pytorch程序(指定单/多显卡)(用gpu运行python),希望有所帮助,仅作参考,欢迎阅读内容。
文章相关热门搜索词:在gpu上运行的程序,pytorch gpu不可用,pytorch gpu运行,pytorch gpu运行,在GPU上运行MATLAB程序,gpu版本的pytorch可以用cpu跑吗,如何在gpu上运行python,如何在gpu上运行python,内容如对您有帮助,希望把文章链接给更多的朋友!
torch.cuda常用指令https://blog.csdn.net/qq_43307074/article/details/127628498?spm=1001.2014.3001.5501方法1和方法2可以同时使用,比如在运行代码时使用:
CUDA_VISIBLE_DEVICES=2,3 python test.py而在代码内部又指定:
model.cuda(1)那么model会在GPU3上运行。原理是CUDA_VISIBLE_DEVICES遍历当前可见的设备,并从零开始为可见设备编号。CUDA_VISIBLE_DEVICES使得只有GPU2,3可见,程序就会把这两张显卡编号为GPU0,1,2,3,cuda(1)把model加载到了GPU1上,则实际使用的显卡是GPU3。
如果利用.cuda()或torch.cuda.set_device()把模型加载到多个显卡上,而实际上只使用一张显卡运行程序的话,那么程序会把模型加载到第一显卡上,在运行代码时使用:
CUDA_VISIBLE_DEVICES=2,3 python test.py而在代码内部又指定:
model.cuda('cuda:1,0')那么model会在GPU3上运行。
3. 使用device3.1. 检查所用的deviceimport torchx = torch.tensor([[1,2,3],[4,5,6]])print(x.device)3.2. 模型/数据指定GPU,具体有以下几种形式:import torchcuda = torch.device('cuda:0'/'cuda')x = torch.tensor([1,2,3],device = cuda)print(x.device)import torchdevice = torch.device('cuda:0'/'cuda')x = torch.rand((4,5)).to(device)print(x.device)import torchx = torch.tensor([1,2,3],device = torch.device('cuda:0'/'cuda')print(x.device)import torchdevice = torch.device('cuda:0'/'cuda')...net = net.to(device)print(net.device)经典写法:
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")如果想要在其他GPU上运行,举例如下:
import torchimport os os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'cuda = torch.device('cuda:1')x = torch.tensor([[1,2,3]],device = cuda)print(x.device)4. 使用torch.nn.DataParallel多卡数据并行一般使用torch.nn.DataParallel
torch.nn.DataParallel(model,device_ids)举例:torch.nn.DataParallel(model, device_ids=device_ids)torch.nn.DataParallel(modul, device_ids=[x1,x2,x3,x4,x5,x6,x7])# 使用的GPU一定是编号连续的torch.nn.DataParallel(model,device_ids = range(torch.cuda.device_count()) )其中model是需要运行的模型,device_ids指定部署模型的显卡,数据类型是list/device。
device_ids中的第一个GPU(即device_ids[0])和model.cuda()或torch.cuda.set_device()中的第一个GPU序号应保持一致,否则会报错。此外如果两者的第一个GPU序号都不是0,比如设置为:
model=torch.nn.DataParallel(model,device_ids=[2,3])model.cuda(2)那么程序可以在GPU2和GPU3上正常运行。device_ids的默认值是使用可见的GPU,不设置model.cuda()或torch.cuda.set_device()等效于设置了model.cuda(0)
具体举例如下:
model = model.cuda()device_ids = [0, 1] model = torch.nn.DataParallel(model, device_ids=device_ids)model = model.to(device)model = nn.DataParallel(model, device_ids=[0,1,2,3])device_ids = [0, 1]model = torch.nn.DataParallel(model, device_ids=device_ids).cuda()所有的数据要先load到主GPU上,然后再分发给每个GPU去train,因此应当将model先放到GPU上,然后在进行并行训练。
参考在pytorch中指定显卡https://zhuanlan.zhihu.com/p/166161217
上一篇:vue查询数据el-table不更新数据(vue 查询)
下一篇:威斯康星州的白尾母鹿和小鹿 (© Karel Bock/Shutterstock)(威斯康星州1976)
友情链接: 武汉网站建设