pytorch 多GPU 訓練

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1, 2'
import torch  

#注意以上兩行先后順序不可弄錯

 

device = torch.device('cuda')

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

 

這樣模型就會在gpu 0, 1, 2 上進行訓練