运用pytorch轻松搞定模型单机多GPU并行训练

使用单GPU进行模型的训练时会出现CUDA RuntimeError：内存不足的错误，此时你可以降低batch_size或是图像的大小来进行重新的训练尝试，但是假如你有多块GPU则可以进行多GPU的训练。

本人习惯使用pytorch这个框架，因此列出代码供参考。

import torch.nn as nn

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = MyModel.Model()
model = torch.nn.DataParallel(model, device_ids=[0,1]) #multi_GPU，两块显卡，3块为device_ids=[0,1，3]，以此类推
torch.save(model.module.state_dict(), file_name)  #单GPU为torch.save(model.state_dict(), file_name)

但DataParallel存在GPU使用不均衡的问题，下图给出了相应解释：

想要解决GPU使用不均衡的问题需要将每部分输出都保留在原有的GPU上，而不汇集于GPU-1。目前有两种方法大家可以尝试：

1. 张航开源了名为PyTorch-Encoding的包，好好查询一下用法。
2. 控制多个服务器进行分布式处理。

单机多GPU训练的模型在调用中与单GPU有相应的区别，将在下一讲《多GPU训练模型的调用------pytorch》中进行说明。