Safetensors 怎么只读取指定权重并查看文件元数据

Safetensors 不需要把整个检查点一次性读进内存，才能知道里面有什么。更稳妥的做法是先用 safe_open() 打开文件，读取 f.keys() 得到权重名称，用 f.metadata() 查看文件级字符串元数据，确认目标键存在后再调用 f.get_tensor(name) 取单个张量。

官方地址：https://huggingface.co/docs/safetensors/

下面的示例以 Python 的 PyTorch 绑定为例，重点解决三个容易混淆的问题：文件元数据和权重元数据不是一回事；keys() 只是名称索引；只有调用具体读取方法时，才真正需要目标张量的数据。

先把文件级元数据和权重键分开

Safetensors 的文件头包含每个张量的名称、dtype、shape 和数据偏移，也可以带一个名为 __metadata__ 的字符串键值表。Python 绑定把这两层信息分别暴露为 keys() 和 metadata()。因此，排查文件时不要把“看到了一个键”误认为“已经读取了这个权重”。

from safetensors import safe_open

file_path = "model.safetensors"

# 只打开 Safetensors 容器，设备先固定为 CPU，便于做结构检查。
with safe_open(file_path, framework="pt", device="cpu") as handle:
    tensor_names = list(handle.keys())
    file_metadata = handle.metadata()

    # 文件可能没有 __metadata__，所以用 None 作为正常分支处理。
    print("权重数量:", len(tensor_names))
    print("前几个权重:", tensor_names[:5])
    print("文件元数据:", file_metadata or {})

metadata() 返回的是文件级的可选字符串映射，可能是 None；它不会替代某个张量自己的形状和类型信息。要查看单个张量的 shape 和 dtype，应该针对目标键使用切片对象或读取后的张量。

确认目标键后只读取一个权重

模型权重名称通常很长，直接手写容易漏掉层级或下划线。先列出键，再用集合表达真正需要的目标名称；如果名称不存在，立即给出候选键，比直接调用 get_tensor() 后再猜原因更容易定位。

from safetensors import safe_open

file_path = "model.safetensors"
requested = "encoder.layer.0.attention.self.query.weight"

with safe_open(file_path, framework="pt", device="cpu") as handle:
    # 读取名称索引，不把所有权重值都 materialize 到内存。
    available = set(handle.keys())
    if requested not in available:
        # 用前缀帮助定位命名差异，避免把缺键误判成文件损坏。
        prefix = requested.rsplit(".", 1)[0]
        suggestions = sorted(name for name in available if name.startswith(prefix))
        raise KeyError(f"找不到权重 {requested!r}，相近键: {suggestions[:5]}")

    # 只取请求的一个张量；device 决定张量落在哪个设备。
    selected = handle.get_tensor(requested)
    print("目标键:", requested)
    print("shape:", tuple(selected.shape))
    print("dtype:", selected.dtype)

这里的关键是 get_tensor(requested) 只接收一个权重名。它适合需要实际数值的场景，例如检查某一层、计算摘要或把一个小张量传给后续逻辑；如果只是查看结构，没必要对所有键循环调用它。

只看形状和类型时先用切片接口

如果目标只是判断权重形状是否匹配，先取得切片视图通常更节省。get_slice(name) 返回的是该键的惰性视图，可以读取 get_shape() 和 get_dtype()；确认形状后，再决定是否需要 get_tensor(name)。

from safetensors import safe_open

file_path = "model.safetensors"
target = "decoder.embed_tokens.weight"

with safe_open(file_path, framework="pt", device="cpu") as handle:
    # 先检查键是否存在，避免把命名错误传给切片接口。
    if target not in handle.keys():
        raise KeyError(f"文件中不存在目标权重: {target}")

    view = handle.get_slice(target)
    shape = view.get_shape()
    dtype = view.get_dtype()
    print("结构信息:", {"name": target, "shape": shape, "dtype": dtype})

    # 只有确实需要数值时才加载完整目标张量。
    # tensor = handle.get_tensor(target)

在这条路径里，“只读取指定权重”有两层含义：第一层是按键选择，不遍历并加载所有权重；第二层是先读结构信息，避免为了判断 shape 而立刻创建完整张量。最终是否使用切片或完整读取，取决于后续任务。

四类现象分别怎么处理

对于远程 Hugging Face Hub 模型，如果需求只是查看权重列表、dtype、shape 或文件元数据，可以使用官方的 Safetensors 元数据解析能力；它通过读取文件头来减少不必要的权重下载。若需要真实张量数值，仍然要按目标文件和目标键执行实际读取。

一个可复用的最小检查函数

把键检查、元数据读取和可选加载收在一个函数里，调用方就能明确区分“结构检查”和“数据读取”。

from safetensors import safe_open


def inspect_weight(file_path: str, name: str, load_value: bool = False):
    """检查一个权重；按需加载值，避免默认读取整个文件。"""
    with safe_open(file_path, framework="pt", device="cpu") as handle:
        # 在读取前固定名称索引和文件级元数据，便于记录排查上下文。
        names = set(handle.keys())
        metadata = handle.metadata() or {}
        if name not in names:
            return {"found": False, "metadata": metadata, "available": sorted(names)}

        # 结构信息优先从切片视图获得；需要数值时再调用 get_tensor。
        view = handle.get_slice(name)
        result = {
            "found": True,
            "metadata": metadata,
            "shape": view.get_shape(),
            "dtype": view.get_dtype(),
        }
        if load_value:
            result["tensor"] = handle.get_tensor(name)
        return result

这个函数不负责判断模型是否适合训练，也不负责把多个分片合并成新文件；它只回答“文件里有没有这个键、文件写了哪些元数据、目标权重是什么结构，以及调用方是否真的要求数值”。这条边界能让排查日志更清楚，也能减少不必要的显存或内存占用。

常见问题

Safetensors 的 metadata 能保存任意 JSON 吗？

文件格式中的 __metadata__ 是字符串到字符串的映射，不应把嵌套对象或数组直接当作值写入。需要更复杂的配置时，可以把它序列化成字符串，但读取后仍要由应用自行解析。

只调用 keys() 会把所有权重读进内存吗？

它主要用于取得名称索引。真正需要张量内容时再调用指定键的 get_tensor()；如果只需要 shape 和 dtype，则优先使用切片视图。

为什么从分片模型的某个文件里找不到权重？

分片模型的权重名称和文件名由索引映射。先读取 model.safetensors.index.json 中对应的 weight_map，再打开映射到的分片文件，不能假设每个分片都包含所有键。

小结

排查 Safetensors 时可以固定成一条顺序：用 keys() 确认名称，用 metadata() 读取文件级说明，用 get_slice() 先看结构，最后只对确定的目标键调用 get_tensor()。这样既能保留文件的原始边界，也不会为了查看一个权重而无条件加载整份检查点。