PyTorch学习（5）：并行训练模型权重的本地化与加载

1. 并行训练与非并行训练

        在训练深度神经网络时，我们一般会采用CPU或GPU来完成。得益于开源传统，许多算法都提供了完整的开源代码工程，便于学习和使用。随着GPU的普及，GPGPU已经占据了大部分的训练场景。

        我们在这里仅以GPU训练场景做一些说明。

        当我们使用单GPU训练时，我们称之为非并行训练。使用多个GPU共同训练时，我们称之为并行训练。

        当然，绝大多数场景下，我们都是使用一个GPU运行一个训练进程。

2. 模型本地化

（1）CPU训练模型权重本地化

        使用CPU进行训练，之后保存模型或权重文件，本身没有什么特殊之处，参考如下代码即可。

torch.save(model.state_dict(), 'model.pth')

或

torch.save(model, 'model.pth')

        前者只保存了模型权重，后者保存了整个模型，包括计算图、权重以及各种状态等。

（2）单GPU训练模型权重本地化

        和使用单CPU进行训练是一样的，参考上述代码。

（3）多GPU训练模型权重本地化

        在执行多GPU训练时，我们通常会使用DP（torch.nn.DataParallel）或DDP（torch.nn.DistributedDataParallel）的方式，这样我们在训练过程中，模型的权重参数状态字典中是含有“.module”前缀的，因此在保存时，需要妥善处理。

        如果我们虽然使用多个GPU训练，但可以仅保存一个GPU上的权重参数（因为每个GPU上的数据是一样的），那么我们可以采用如下方式。

torch.save(model.module.state_dict(), 'model.pth')

或

torch.save(model.module, 'model.pth')

        如果我们只是保存当前模型，后续Finetune时，整个训练配置是一样的，那么我们可以直接保存。

torch.save(model.state_dict(), 'model.pth')

或

torch.save(model, 'model.pth')

3. 模型加载

        在讲模型加载前，我们提出一个建议，那就是在保存模型训练权重时，我们尽量保存单GPU模型。这之后，无论是在CPU加载还是GPU加载，无论是单GPU加载还是多GPU加载，都能够灵活应付。

（1）单GPU加载

        这种情况下非常简单，只需要注意torch.load时，可以使用map_location指定加载位置即可。

如果保存的是state_dict，那么加载后可以直接通过torch.load_state_dict进行加载。参考代码如下。

device = torch.device("cuda")

model = Model().to(device)

ckpt = torch.load("model.pth", map_location=device)

model.load_state_dict(ckpt)

        如果保存的是完整的model，参考如下代码。

device = torch.device("cuda")

model = Model().to(device)

ckpt = torch.load("model.pth", map_location=device)

model.load_state_dict(ckpt.state_dict())

（2）多GPU加载

        多GPU加载只需要注意一个点，那就是我们尽量先加载权重再执行DP或DDP操作。

device = torch.device("cuda")

model = Model().to(device)

ckpt = torch.load("model.pth", map_location=device)

model.load_state_dict(ckpt)

model = torch.nn.DataParallel(model)

        如果保存的是完整的model，参考如下代码。

device = torch.device("cuda")

model = Model().to(device)

ckpt = torch.load("model.pth", map_location=device)

model.load_state_dict(ckpt.state_dict())

model = torch.nn.DataParallel(model)