1. 程式人生 > >pytorch 多GPU訓練

pytorch 多GPU訓練

當一臺伺服器有多張GPU時,執行程式預設在一張GPU上執行。通過多GPU訓練,可以增大batchsize,加快訓練速度。

from torch.nn import DataParallel
num_gpu = torch.cuda.device_count()
net = DataParallel(net, device_ides=range(num_gpu))

device_ides預設情況下是所有GPU顯示卡,也可以自己指定用哪幾張卡