Pytorch DataParallel() 多GPU训练

1 目的

将网络中的参数放入到多GPU 中进行训练

2 torch.nn.parallel.DataParallel(module, device_ids, output_device, dim)

参数:

module: 需要进行多 GPU 训练的参数

device_ids: GPU编号( 注意是列表)

另外两个参数都设置成默认就好

3 举例

import toech
import torch.nn as nn

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

class AlexNet(nn.Module):
    ...

if torch.cuda.device_count() > 1:
    alexnet_muti_GPU = nn.parallel.DataPallel(alexnet, device_ids = [0, 1, 2, 3])
alexnet_muti_GPU.to(device)    

注意:

• device_ids 是一个列表

• 需要对 DataParallel() 后的网络赋予一个新的网络 (有点像重新定义)

• 记得对重新定义后的网络进行 .to(device) 放入 GPU 操作