PyTorch 模型保存与加载示例

以下是一个保存 PyTorch 模型的示例代码：

import torch
import torch.nn as nn

# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(10, 1)

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

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

# 定义保存路径
save_path = 'model.pth'

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

在这个例子中，我们首先定义了一个简单的模型MyModel，它包含一个线性层。然后，我们创建了一个模型实例model。接下来，我们定义了保存路径save_path，这是我们将模型保存到的文件路径。最后，我们使用torch.save()函数将模型的状态字典保存到指定路径。

要加载已保存的模型，可以使用以下代码：

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

# 定义加载路径
load_path = 'model.pth'

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

在这个例子中，我们首先创建了一个新的模型实例model。然后，我们定义了加载路径load_path，这是我们从中加载模型的文件路径。最后，我们使用model.load_state_dict()函数加载保存的模型的状态字典。