pytorch multi-gpu train

Posted zf-blog

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了pytorch multi-gpu train相关的知识,希望对你有一定的参考价值。

记录一下pytorch如何进行单机多卡训练:

官网例程:https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html

下面以一个例子讲解一下,例如现在总共有8张卡,在第5、6、7三张卡上进行训练;

step 1:可视化需要用到的GPU

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "5 , 6 , 7"

device = torch.device("cuda:0")  #注意多卡训练的时候,默认都先将model和data先保存在id:0的卡上(即实际的第5块卡),然后model的参数会复制共享到其他卡上,data也会平分成若干个batch到其他卡上(所以第一块卡上稍微耗费一点显存);

device_ids = [0 , 1 , 2] #注意device_ids必须从0开始,代码中的所有的device id都需要从0开始(这里的0代表第5块卡,1代表第6块卡,类推);

step 2:利用DataParallel对Model类进行封装

model = nn.DataParallel(model , device_ids = device_ids)

model.to(device)

step 3:

data.to(device)  #id:0卡上的数据再被平分成若干个batch到其他卡上

注意:晚上还有一些例程,需要对optimizer和loss利用DataParellel进行封装,没有试验过,但上面方法是参考官网例程,并经过实操考验;

 

以上是关于pytorch multi-gpu train的主要内容,如果未能解决你的问题,请参考以下文章

pytorch设置多GPU运行的方法

PyTorch学习笔记:PyTorch生态简介

PyTorch学习笔记:PyTorch进阶训练技巧

pytorch 中的常用矩阵操作

Pytorch Note1 Pytorch介绍

pytorch_geometric + MinkowskiEngine