Pytorch：多GPU训练网络与单GPU训练网络保存模型的区别

分类专栏： PyTorch

• 测试环境：Python3.6 + Pytorch1.1

在pytorch中，使用多GPU训练网络需要用到 【nn.DataParallel】：

1. gpu_ids = [0, 1, 2, 3]

2. device = t.device("cuda:0" if t.cuda.is_available() else "cpu") # 只能单GPU运行

3. net = LeNet()

4. if len(gpu_ids) > 1:

5. net = nn.DataParallel(net, device_ids=gpu_ids)

6. net = net.to(device)

而使用单GPU训练网络：

1. device = t.device("cuda:0" if t.cuda.is_available() else "cpu") # 只能单GPU运行

2. net = LeNet().to(device)

由于多GPU训练使用了 nn.DataParallel(net, device_ids=gpu_ids) 对网络进行封装，因此在原始网络结构中添加了一层module。网络结构如下：

1. DataParallel(

2. (module): LeNet(

3. (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))

4. (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))

5. (fc1): Linear(in_features=400, out_features=120, bias=True)

6. (fc2): Linear(in_features=120, out_features=84, bias=True)

7. (fc3): Linear(in_features=84, out_features=10, bias=True)

8. )

9. )

而不使用多GPU训练的网络结构如下：

1. LeNet(

2. (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))

3. (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))

4. (fc1): Linear(in_features=400, out_features=120, bias=True)

5. (fc2): Linear(in_features=120, out_features=84, bias=True)

6. (fc3): Linear(in_features=84, out_features=10, bias=True)

7. )

由于在测试模型时不需要用到多GPU测试，因此在保存模型时应该把module层去掉。如下：

1. if len(gpu_ids) > 1:

2. t.save(net.module.state_dict(), "model.pth")

3. else:

4. t.save(net.state_dict(), "model.pth")

 

参考 

• https://blog.youkuaiyun.com/tsq292978891/article/details/83586935