pytorch多GPU并行训练网络模型

训练网络模型无非分为几步：

1. 加载数据集;
2. 初始化网络模型;
3. 初始化优化函数和定义学习率;
4. 初始化模型参数;
5. 保存模型参数;

我们采用多卡并行训练网络模型时，要注意第一点就是加载数据的batch_size = 原先batch_size * GPU个数，且batch_size不能小于GPU的个数。
首先要检测我们的服务器有几张显卡：

USE_MULTI_GPU = True

# 检测机器是否有多张显卡
if USE_MULTI_GPU and torch.cuda.device_count() > 1:
    MULTI_GPU = True
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1"
    device_ids = [0, 1]
else:
    MULTI_GPU = False
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

其中os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1"是将机器中的GPU进行编号。
接下来就是读取模型：

net = Vgg19()
if MULTI_GPU:
    net = nn.DataParallel(net,device_ids=device_ids)
net.to(device)

这里与单卡的区别就是多了nn.DataParallel这一步操作。
接下来是optimizer和scheduler的定义：

optimizer=optim.Adam(net.parameters(), lr=1e-3)
scheduler = StepLR(optimizer, step_size=100, gamma=0.1)
if MULTI_GPU:
    optimizer = nn.DataParallel(optimizer, device_ids=device_ids)
    scheduler = nn.DataParallel(scheduler, device_ids=device_ids)
    optimizer = optimizer.module
    scheduler = scheduler.module

训练模型同时还需要保存模型：

torch.save(net.module.state_dict(), save_path)

与单卡不同的就是加了module。

一、这边要注意如果多卡加载单卡训练的预训练模型：

structure_model = torch.load('./latest_net_s_gen.pth')
if isinstance(self.s_gen, nn.DataParallel):
    self.s_gen = self.s_gen.module
self.s_gen.load_state_dict(structure_model, False)

二、多卡训练模型加载多卡预训练模型:

self.netG.load_state_dict(torch.load(path))

三、单卡训练模型加载单卡预训练模型:

self.netG.load_state_dict(torch.load(path))

四、单卡训练模型加载多卡预训练模型:

state_dict = torch.load('checkpoint.pth’)
from collections import OrderedDict
new_state_dict = OrderedDict()
for k,v in state_dict.items():
	name = k[7:]
	new_state_dict[name]  =v 
self.netG.load_state_dict(new_state_dict)

参考：
https://blog.csdn.net/weixin_43402775/article/details/108887826
https://www.cnblogs.com/c-chenbin/p/14809822.html
测试机子是否可以采用多卡训练，可以采用LeNet代码进行测试：https://github.com/CHENBIN99/LeNet-5_by_Pytorch