如何在 PyTorch 中保存和加载模型

PyTorch模型保存与加载

• 一、保存模型参数

  在 PyTorch 中，保存模型参数是最常用的方法。使用 torch.save() 函数将模型的 state_dict 保存到文件中。state_dict 是一个 Python 字典对象，包含了模型的所有参数。保存时可以指定文件路径和文件名。例如：

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

  这种方法的优点是文件较小，适合在相同代码环境下恢复模型。加载时需要先定义模型结构，再调用 load_state_dict() 方法加载参数。例如：

  model = TheModelClass(*args, **kwargs)

  model.load_state_dict(torch.load('model_params.pth'))

  model.eval() # 切换到评估模式

  这种方法适合大部分场景，尤其是需要在不同机器上重复实验时。

• 二、保存整个模型

  保存整个模型包括模型结构和参数，这种方法使用 torch.save() 直接保存模型对象。例如：

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

  加载时使用 torch.load() 函数直接恢复模型：

  model = torch.load('model.pth')

  model.eval() # 切换到评估模式

  这种方法的优点是简单，文件包含所有信息，适合在不熟悉代码结构的情况下恢复模型。缺点是文件较大，不适合频繁更新模型。

• 三、保存优化器状态

  训练过程中，优化器的状态同样重要，需要保存以便恢复训练。保存优化器状态的方法类似于保存模型参数。使用 torch.save() 函数保存优化器的 state_dict，例如：

  torch.save(optimizer.state_dict(), 'optimizer_params.pth')

  加载时需要先定义优化器，再调用 load_state_dict() 方法加载参数：

  optimizer.load_state_dict(torch.load('optimizer_params.pth'))

  这样可以确保训练过程中的动量和学习率等参数得到完全恢复。

• 四、保存和加载自定义对象

  在深度学习项目中，经常需要保存和加载自定义的类和对象。可以使用 Python 的序列化工具 pickle 进行操作。例如，将自定义对象保存到文件：

  import pickle

  with open('custom_obj.pkl', 'wb') as f:

  pickle.dump(custom_obj, f)

  加载时使用：

  with open('custom_obj.pkl', 'rb') as f:

  custom_obj = pickle.load(f)

  自定义对象的类定义必须在加载环境中存在，否则会报错。这种方法适用于保存数据处理管道、训练过程中的状态等自定义内容。