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)