Pytorch单、多GPU和CPU训练模型保存和加载

Pytorch多GPU训练模型保存和加载

在多GPU训练中，模型通常被包装在torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel中，这会在模型的参数名前加上module前缀。因此，在保存模型时，需要使用model.module.state_dict()来获取模型的状态字典，以确保保存的参数名与模型定义中的参数名一致。（本质上原来的model还是存在的，参数也会同步更新）

1. 多GPU训练模型保存

   在多GPU训练时，模型通常被包装在torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel中，这会在模型的参数名前加上module前缀。因此，在保存模型时，需要使用model.module.state_dict()来获取模型的状态字典，以确保保存的参数名与模型定义中的参数名一致。

2. 单GPU或CPU加载模型

   当在单GPU或CPU上加载模型时，如果直接使用model.state_dict()保存的模型，由于缺少module前缀，会导致参数名不匹配，从而无法正确加载模型。因此，在保存多GPU训练的模型时，应该使用model.module.state_dict()来保存模型的状态字典，这样在单GPU或CPU上加载模型时，可以直接加载，不会出现参数名不匹配的问题。

3. 示例代码

   以下是一个示例代码，展示了如何在多GPU训练时保存模型，并在单GPU或CPU上加载模型：

import torch
import torch.nn as nn
import os
os.os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"	#设置GPU编号
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 假设这是你的模型定义
class YourModel(nn.Module):
    def __init__(self):
        super(YourModel, self).__init__()
        self.fc = nn.Linear(10, 2)

    def forward(self, x):
        return self.fc(x)

# 创建模型实例
model = YourModel()

# 将模型移动到多GPU上
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)
    model = model.to(device)
else:
	model = model.to(device)
······
# 假设这是你的训练代码，训练完成后保存模型
if torch.cuda.device_count() > 1:
    torch.save(model.module.state_dict(), 'model.pth')
else:
    torch.save(model.state_dict(), 'model.pth')

# 在单、多GPU或CPU上加载模型
model = YourModel()
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
model.load_state_dict(torch.load('model.pth'))
model = model.to(device)

2 在多GPU训练得到的模型加载时，通常需要考虑以下几个步骤：

1. 模型保存
   在多GPU训练时，模型通常被包装在torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel中。因此，在保存模型时，需要确保保存的是模型的state_dict而不是整个模型对象。例如：

if torch.cuda.device_count() > 1:
    torch.save(model.module.state_dict(), 'model.pth')
else:
    torch.save(model.state_dict(), 'model.pth')

2. 模型加载
   在加载模型时，首先需要创建模型的实例，然后使用load_state_dict方法来加载保存的权重。如果模型是在多GPU环境下训练的，那么在加载时也应该使用torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel来包装模型。例如：

model = YourModel()
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
model.load_state_dict(torch.load('model.pth'))
model = model.to('cuda')

3. 注意事项
   在加载模型时，需要注意以下几点：

如果模型是在多GPU环境下训练的，那么在加载时也应该使用相同数量的GPU，或者使用torch.nn.DataParallel来包装模型，即使只有一个GPU可用。

如果模型是在分布式训练环境下训练的，那么在加载时也应该使用torch.nn.parallel.DistributedDataParallel来包装模型。

如果模型是在混合精度训练（如使用了torch.cuda.amp）下训练的，那么在加载模型后，应该恢复之前的精度设置。

3 为了避免模型保存和加载出错

在多GPU训练的模型使用了torch.nn.DataParallel来包装模型，但本质上原来的model是依然存在的，且参数会同步更新：

1. torch.nn.DataParallel 的工作原理
   torch.nn.DataParallel 是 PyTorch 提供的一个类，用于在多个 GPU 上并行训练模型。它的工作原理如下：
   模型复制：DataParallel 会在每个 GPU 上创建模型的副本。
   数据分发：输入数据会被分发到各个 GPU 上。
   前向传播：每个 GPU 上的模型副本会独立进行前向传播计算。
   梯度收集：所有 GPU 上的梯度会被收集并汇总到主 GPU 上。
   参数更新：主 GPU 上的优化器会根据汇总后的梯度更新模型参数，然后将更新后的参数同步回其他 GPU。
2. 模型参数更新
   当你使用 model_train = torch.nn.DataParallel(model) 后，model_train 实际上是一个包装了原始模型 model 的对象。虽然 model_train 是多GPU并行的版本，但它的参数更新是通过主 GPU 上的优化器完成的，并且这些更新会同步回原始模型 model。
   因此，model 的参数确实会被更新。具体来说：
   前向传播和反向传播：在 train_model 函数中，model_train 用于前向传播和反向传播。
   参数更新：优化器 optimizer 使用的是 model.parameters()，即原始模型的参数。在每次迭代中，优化器会根据汇总后的梯度更新这些参数。
   参数同步：更新后的参数会自动同步到 model_train 中的各个 GPU 副本。
   因此可以使用如下代码，加载模型和保存模型：

import torch
import torch.nn as nn
import os
os.os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"	#设置GPU编号
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 假设这是你的模型定义
class YourModel(nn.Module):
    def __init__(self):
        super(YourModel, self).__init__()
        self.fc = nn.Linear(10, 2)

    def forward(self, x):
        return self.fc(x)

# 创建模型实例
model = YourModel()

# 将模型移动到多GPU上，单GPU依然适用
if torch.cuda.device_count() > 1:
	model_train = nn.DataParallel(model)
	model_train = model_train.to(device)
else:
	model_train = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)#注意这是model的参数
······
output = model_train(input)	# 多卡时训练的输入和输出，注意这是model_train

# 假设这是你的训练代码，训练完成后保存模型
torch.save(model.state_dict(), 'model.pth')	#注意这是model

• 再在单/多GPU或CPU上加载模型，都不会报错，因为这里的model不是包装体，不带module

model = YourModel()
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
model.load_state_dict(torch.load('model.pth',map_location = device))
model = model.to(device)