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 操作