加载.npy到PyTorch时遇键不匹配及形状不匹配错误求助

1. 键不匹配问题 #

• 核心原因：.npy文件的参数键名和PyTorch版S3FD的state_dict键名规则不一致。比如原参数可能是Caffe导出的扁平键名（如conv1_1.weight），而PyTorch模型的键名带层级前缀（如base.conv1_1.weight）；或是多GPU训练的模型state_dict带module.前缀，和单GPU模型键名不兼容。
• 解决方法：
  • 先打印双方键列表对比：

    # 打印PyTorch模型的键
    print([k for k in model.state_dict().keys()])
    # 打印.npy加载后的键（假设.npy是字典格式）
    npy_params = np.load('params.npy', allow_pickle=True).item()
    print([k for k in npy_params.keys()])
    

  • 编写键名映射函数，比如把conv1_1.weight映射为base.conv1_1.weight，注意不要遗漏偏置项（bias）。

2. 参数类型错误 #

• 核心原因：直接将numpy数组赋值给state_dict，未转换为torch.Tensor；或是转换时未匹配模型的 dtype（S3FD默认用float32）。
• 解决方法：
  • 对每个numpy参数执行张量转换：

    # 转换为float32张量，若模型在GPU上需同步设备
    torch_param = torch.from_numpy(npy_arr).float().to(model.device)
    

  • 确保所有参数都完成转换后再更新state_dict，禁止混合numpy数组和张量。

3. 参数形状不匹配问题 #

• 常见原因：
  • 框架间维度顺序差异：比如部分自定义层的权重维度顺序和PyTorch默认规则不符；
  • 批量归一化层参数不对应：PyTorch的BN层有weight/bias/running_mean/running_var四个参数，而.npy可能只保存了前两个，或顺序混乱；
  • .npy参数被扁平化存储，加载后未恢复原形状。
• 解决方法：
  • 逐个对比参数形状：

    torch_shape = model.state_dict()[torch_key].shape
    npy_shape = npy_params[npy_key].shape
    if torch_shape != npy_shape:
        print(f"不匹配：{torch_key} 预期{torch_shape}，实际{npy_shape}")
    

  • 针对差异调整：比如卷积权重转置、全连接层维度重排；BN层需手动补充缺失的running_mean/running_var（初始化为0和1）。

通用排查流程 #

先加载空模型的state_dict，逐个匹配.npy参数，分步验证：

import torch
import numpy as np
from your_model_module import S3FD

model = S3FD()
model_state = model.state_dict()
npy_params = np.load('s3fd_params.npy', allow_pickle=True).item()

# 自定义键映射函数，根据实际情况调整
def map_key(torch_key):
    return torch_key.replace('base.', '')  # 示例：移除PyTorch模型的层级前缀

for torch_key, param in model_state.items():
    npy_key = map_key(torch_key)
    if npy_key not in npy_params:
        print(f"缺失键：{npy_key}（对应PyTorch键：{torch_key}）")
        continue
    npy_arr = npy_params[npy_key]
    # 形状检查与调整
    if param.shape != npy_arr.shape:
        # 示例：尝试转置卷积权重
        npy_arr = npy_arr.transpose(1, 0, 2, 3)
    # 转换张量并赋值
    model_state[torch_key] = torch.from_numpy(npy_arr).float()

# 先宽松加载，查看剩余不匹配项
model.load_state_dict(model_state, strict=False)

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