PyTorch 中如何保存和加载模型的 state_dict？

推荐答案

在 PyTorch 中，保存和加载模型的 state_dict 是常见的操作。state_dict 是一个 Python 字典对象，它将每一层映射到其对应的参数张量。以下是保存和加载 state_dict 的推荐方法：

保存模型的 state_dict

import torch

# 假设 model 是已经训练好的模型
torch.save(model.state_dict(), 'model_state_dict.pth')

加载模型的 state_dict

// www.javascriptcn.com code example
import torch
from model_definition import MyModel  # 假设 MyModel 是模型的定义类

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

# 加载 state_dict
model.load_state_dict(torch.load('model_state_dict.pth'))

# 将模型设置为评估模式
model.eval()

本题详细解读

1. 保存模型的 state_dict

• model.state_dict(): 这个方法返回一个包含模型所有参数（权重和偏置）的字典。字典的键是参数的名称，值是参数的张量。
• torch.save(): 这个方法将 state_dict 保存到指定的文件中。文件通常以 .pth 或 .pt 作为扩展名。

2. 加载模型的 state_dict

• torch.load(): 这个方法从文件中加载保存的 state_dict。
• model.load_state_dict(): 这个方法将加载的 state_dict 应用到模型中，恢复模型的参数。
• model.eval(): 在加载模型后，通常需要将模型设置为评估模式（eval()），这会影响到某些层的行为，如 Dropout 和 BatchNorm。

3. 注意事项

• 模型定义: 在加载 state_dict 之前，必须确保模型的定义与保存 state_dict 时的模型定义一致。否则，加载时可能会出错。

• 设备: 如果模型是在 GPU 上训练的，而加载时在 CPU 上，可能需要使用 map_location 参数来指定设备：

  model.load_state_dict(torch.load('model_state_dict.pth', map_location=torch.device('cpu')))

• 优化器状态: 如果还需要保存和加载优化器的状态，可以类似地保存和加载优化器的 state_dict：

  torch.save(optimizer.state_dict(), 'optimizer_state_dict.pth')
  optimizer.load_state_dict(torch.load('optimizer_state_dict.pth'))

通过以上方法，可以有效地保存和加载 PyTorch 模型的 state_dict，从而在训练中断后恢复训练，或者在不同的环境中使用相同的模型参数。