其他分享
首页 > 其他分享> > pytorch multi-gpu train

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进行封装,没有试验过,但上面方法是参考官网例程,并经过实操考验;

 

标签:例程,卡上,ids,pytorch,train,device,gpu,model,data
来源: https://www.cnblogs.com/zf-blog/p/10599328.html