PyTorch加载权重时提示“Weights only load failed”，如何解决？

Question

在使用PyTorch加载模型权重时，遇到“Weights only load failed”错误，可能是由于模型结构与权重不匹配、权重文件损坏或路径错误导致。如何系统性地排查并解决这一问题，确保模型架构与预训练权重完全对应，同时验证权重文件的完整性和正确性？

Answer 1

1. 初步问题分析与常见原因

在使用PyTorch加载模型权重时，如果遇到“Weights only load failed”错误，通常是由以下三个主要原因导致的：

• 模型结构与权重不匹配：预训练权重与当前模型的架构定义不一致。
• 权重文件损坏：权重文件可能在传输或存储过程中被破坏。
• 路径错误：指定的权重文件路径无效或不可访问。

为了解决这一问题，我们需要从以下几个方面进行系统性排查：

1. 验证模型结构是否与权重完全对应。
2. 检查权重文件的完整性和正确性。
3. 确保路径配置无误。

2. 模型结构一致性检查

模型结构与权重的匹配是成功加载的关键。以下是具体步骤：

# 假设我们有一个预训练权重文件和一个自定义模型
import torch

model = YourModel()  # 定义模型
pretrained_weights = torch.load("path_to_weights.pth")

# 检查模型参数键值是否与权重文件中的键值一致
model_keys = set(model.state_dict().keys())
weights_keys = set(pretrained_weights.keys())

missing_keys = model_keys - weights_keys
unexpected_keys = weights_keys - model_keys

print(f"Missing keys in pretrained weights: {missing_keys}")
print(f"Unexpected keys in pretrained weights: {unexpected_keys}")

通过上述代码，可以明确模型中缺少哪些权重键值以及权重文件中多余的键值。这有助于判断模型结构与权重是否一致。

3. 权重文件完整性验证

为了确保权重文件未被损坏，可以通过以下方法验证其完整性：

方法	描述
MD5校验	计算权重文件的MD5哈希值并与官方提供的值对比。
文件大小检查	确认文件大小是否与预期一致。
尝试加载	通过torch.load加载文件并捕获异常。

例如，使用以下代码验证文件是否可加载：

try:
    weights = torch.load("path_to_weights.pth")
    print("Weights loaded successfully.")
except Exception as e:
    print(f"Error loading weights: {e}")

4. 路径配置排查

路径错误可能导致无法找到权重文件。以下是排查路径问题的步骤：

1. 确保路径拼写正确且文件存在。
2. 检查文件权限，确保程序有读取权限。
3. 对于远程路径，确认网络连接正常。

可以通过以下代码检查路径有效性：

import os

weight_path = "path_to_weights.pth"
if not os.path.exists(weight_path):
    print(f"Path does not exist: {weight_path}")
else:
    print("Path exists.")

5. 系统性排查流程图

以下是完整的排查流程图，帮助用户按步骤解决问题：

graph TD;
    A[开始] --> B{路径是否正确?};
    B -- 是 --> C{模型结构是否匹配?};
    B -- 否 --> D[修正路径];
    C -- 是 --> E{权重文件是否完整?};
    C -- 否 --> F[调整模型结构];
    E -- 是 --> G[加载成功];
    E -- 否 --> H[修复权重文件];