【pytorch】多GPU训练

使用多GPU训练pytorch模型只需要加一句DataParallel即可，如下

from torch.nn import DataParallel os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2" # 允许使用的GPU model = .. (torch.nn.Module) model = DataParallel(model) model.cuda()

数据和标签正常使用

data.cuda() label.cuda()