PyTorch 模型保存与加载教程：简单易懂的步骤

PyTorch 模型保存与加载教程：简单易懂的步骤

本文将为您详细介绍如何在 PyTorch 中保存和加载训练好的模型，并提供清晰的示例代码。

1. 使用 torch.save() 函数保存模型

torch.save() 函数可以将训练好的模型保存到文件中。该函数需要两个参数：要保存的模型和文件名。

示例代码：

import torch

# 定义模型
model = torch.nn.Linear(10, 1)

# 训练模型...

# 保存模型
torch.save(model.state_dict(), 'model.pth')

2. 使用 torch.load() 函数加载模型

torch.load() 函数可以加载保存在文件中的模型。该函数需要一个参数：模型文件的路径。

示例代码：

import torch

# 定义模型
model = torch.nn.Linear(10, 1)

# 加载模型
model.load_state_dict(torch.load('model.pth'))

**注意：**加载模型时，需要先定义模型的结构，然后再加载模型参数。如果模型的结构和保存模型时不同，加载模型时会出现错误。