Mathematica中文论坛-非官方

标题: mulit gpu training [打印本页]

作者: meatball1982 时间: 4 天前
标题: mulit gpu training
https://blog.csdn.net/RuanChengfeng/article/details/128172761

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

model = CreateModel()

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

CUDA_VISIBLE_DEVICES=1,3

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

model = CreateModel()

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

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)

欢迎光临 Mathematica中文论坛-非官方 (http://www.ilovemathematica.com/)