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上加载模型:

python 复制代码
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而不是整个模型对象。例如:
python 复制代码
if torch.cuda.device_count() > 1:
    torch.save(model.module.state_dict(), 'model.pth')
else:
    torch.save(model.state_dict(), 'model.pth')
  1. 模型加载
    在加载模型时,首先需要创建模型的实例,然后使用load_state_dict方法来加载保存的权重。如果模型是在多GPU环境下训练的,那么在加载时也应该使用torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel来包装模型。例如:
python 复制代码
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')
  1. 注意事项
    在加载模型时,需要注意以下几点:

如果模型是在多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 副本。
    因此可以使用如下代码,加载模型和保存模型:
python 复制代码
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
python 复制代码
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)
相关推荐
weixin_599073943 分钟前
第P7周:马铃薯病害识别(VGG-16复现)
python·深度学习·机器学习
好多渔鱼好多10 分钟前
【大模型学习】第二十一章 什么是知识蒸馏
人工智能·机器学习·ai·大模型
自由鬼14 分钟前
数据标注质量对AI模型质量的影响分析
人工智能·深度学习·机器学习·ai
无码不欢的我18 分钟前
在windows10系统上安装docker,然后在容器中运行GPU版本的Pytorch,并使用vscode连接该容器
pytorch·vscode·docker
zhz521418 分钟前
遥感新态势:Sentinel - 2多光谱指数与AI深度融合
人工智能·sentinel·遥感·农业
梦想画家19 分钟前
PyTorch 系列教程:探索自然语言处理应用
人工智能·pytorch·自然语言处理
糖葫芦君31 分钟前
tensorflow与torch并行读取数据机制
人工智能·python·tensorflow
搏博36 分钟前
无人机的飞行路径规划之CH-PPO算法(思考)
人工智能·算法·无人机
小明同学的开发日常1 小时前
【Spring AI】05. 向量数据库-Redis
数据库·人工智能·spring
凤凰AI1 小时前
Pytorch的入门
人工智能·pytorch·python