torch.save(model.state_dict(), ‘best_model.pth‘)

- 日理万妓 2023-10-12 20:05 22阅读 0赞

`torch.save(model.state_dict(), 'best_model.pth')` 是一个将模型的参数字典保存到文件的操作。在这个示例中，`model.state_dict()` 返回了当前模型的参数字典，然后使用 `torch.save()` 函数将参数字典保存到名为 `"best_model.pth"` 的文件中。

这种方式是一种常见的保存模型的方法，它将模型的参数保存为一个二进制文件，以便在需要时进行加载和恢复。保存参数字典而不是整个模型的原因是模型的结构和计算图不需要保存，只需要保存模型的权重和偏置等可学习参数即可。

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

model = MyModel()  # 创建一个与保存模型参数相同的模型实例
    model.load_state_dict(torch.load('best_model.pth'))  # 加载参数字典到模型中

通过这样的加载方式，可以将保存的模型参数加载到一个与原始模型结构相同的模型实例中，从而恢复模型的参数。这样可以在训练过程中保存最佳模型，或在需要时加载模型并进行预测或继续训练。

model.state_dict()

`model.state_dict()` 是一个方法，用于返回模型的参数字典`（state_dict）`。模型的参数字典是一个`Python`字典对象，其中包含了模型的所有可学习参数（权重和偏置）及其对应的张量数值。

在深度学习中，模型的参数通常由各个层的权重和偏置组成。`model.state_dict()` 方法会返回一个字典，其中键是每个参数的名称，值是对应参数的张量数值。

使用 `model.state_dict()` 的主要目的是保存和加载模型的参数。通过调用 `model.state_dict()` 可以获取当前模型的参数字典，然后可以将其保存到文件中。保存模型的参数字典后，可以使用 `torch.load()` 加载字典，并将其加载到模型中，从而恢复模型的参数。

以下是一个示例代码，演示了如何使用 `model.state_dict()` 保存和加载模型的参数：

# 保存模型的参数
    torch.save(model.state_dict(), 'model.pth')
    
    # 加载模型的参数
    model = MyModel()  # 创建一个空白模型
    model.load_state_dict(torch.load('model.pth'))

这样，模型的参数就可以从保存的文件中加载回来，并应用于新的模型实例。这在训练过程中保存和加载模型，或者在不同的会话中迁移模型等场景中非常有用。