Pytorch如何使用多块GPU同时训练

2022-09-23 08:58:21

记录一下:
If you want to use all the available GPUs:

device= torch.device("cuda"if torch.cuda.is_available()else"cpu")

model= CreateModel()

model= nn.DataParallel(model)
model.to(device)

If you want to use specific GPUs: (For example, using 2 out of 4 GPUs)

device= torch.device("cuda:1,3"if torch.cuda.is_available()else"cpu")## specify the GPU id's, GPU id's start from 0.

model= CreateModel()

model= nn.DataParallel(model,device_ids=[1,3])
model.to(device)
  • 作者:xuecaisun
  • 原文链接:https://blog.csdn.net/xuecaisun/article/details/126145760
    更新时间:2022-09-23 08:58:21