加载预训练模型时出现KeyError: 'optimizer_state_dict'问题求助

问题描述 #

使用Lipreading_using_Temporal_Convolutional_Networks仓库的代码，加载该仓库模型 zoo 中的预训练模型，在自有数据集上恢复训练时，运行main.py报错：

Traceback (most recent call last):
  File "main.py", line 348, in <module>
    main()
  File "main.py", line 299, in main
    model, optimizer, epoch_idx, ckpt_dict = load_model(args.model_path, model, optimizer)
  File "/.../lipreading/utils.py", line 170, in load_model
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
KeyError: 'optimizer_state_dict'

原因分析 #

你下载的预训练模型 checkpoint 文件中，只保存了模型的权重参数（model_state_dict），没有保存优化器的状态字典（optimizer_state_dict）。这类公开的预训练模型通常仅提供模型权重，用于微调或推理，不会包含训练过程中的优化器状态。

解决方法 #

1. 确认checkpoint内容 #

先写个简单脚本查看checkpoint里的所有键，确认是否真的没有optimizer_state_dict：

import torch
# 替换成你的模型路径
ckpt_path = "path/to/your/model.pth"
checkpoint = torch.load(ckpt_path, map_location="cpu")
print("Checkpoint包含的键：", list(checkpoint.keys()))

2. 修改加载逻辑 #

根据需求调整utils.py中的load_model函数：

• 如果只需要加载模型权重，重新初始化优化器训练：
  找到utils.py第170行附近的代码，修改为：

  # 加载优化器状态（如果存在）
  if 'optimizer_state_dict' in checkpoint:
      optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
  else:
      print("预训练模型未保存优化器状态，将使用新的优化器配置")
  

  同时，如果checkpoint中没有epoch_idx，也要对应处理起始epoch，比如默认从0开始。

• 若要严格恢复训练状态（需要优化器状态）：
  这种情况需要联系模型提供者获取完整的训练checkpoint，或者放弃恢复优化器状态，从头开始微调模型。

内容的提问来源于stack exchange，提问作者Shema