pytorch指定使用多个GPU

方式一：指定使用所有GPU

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = CreateModel()model= nn.DataParallel(model)model.to(device)

方式二：结合系统环境变量使用特定的GPU

CUDA_VISIBLE_DEVICES=1,3Device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = CreateModel()model= nn.DataParallel(model)model.to(device)

方式三：直接指定特定的GPU

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)

以上所有情况中，后面使用的数据都要映射到对应的GPU设备。

如果x,y是数据：

X.to(device)y.to(device)

来源地址：https://blog.csdn.net/RuanChengfeng/article/details/128172761