第十七章：AI模型的“存档”艺术：揭秘.pt, .ckpt, .safetensors的江湖恩怨

AI保存格式

• 前言：训练的终点，是"存档"的开始
• [第一章：PyTorch的"传统艺能"------ torch.save 与 Pickle 协议](#第一章：PyTorch的“传统艺能”—— torch.save 与 Pickle 协议)
• • [1.1 两种存档策略：只存"灵魂" vs. "灵魂+躯体"一起存](#1.1 两种存档策略：只存“灵魂” vs. “灵魂+躯体”一起存)
  • [1.2 用代码实现两种存档与读档方式](#1.2 用代码实现两种存档与读档方式)
  • [1.3 【高能预警】Pickle的"阿喀琉斯之踵"：恐怖的"代码注入"风险](#1.3 【高能预警】Pickle的“阿喀琉斯之踵”：恐怖的“代码注入”风险)
• [第二章：新王登基 ------ safetensors 的崛起](#第二章：新王登基 —— safetensors 的崛起)
• • [2.1 为什么需要safetensors？------ 安全、快速、跨平台](#2.1 为什么需要safetensors？—— 安全、快速、跨平台)
  • [2.2 safetensors 的核心思想：数据与代码的"彻底分离"](#2.2 safetensors 的核心思想：数据与代码的“彻底分离”)
  • [2.3 用Hugging Face safetensors库进行安全的模型存取](#2.3 用Hugging Face safetensors库进行安全的模型存取)
• [第三章：江湖别传 ------ .ckpt 与 .pth 的故事](#第三章：江湖别传 —— .ckpt 与 .pth 的故事)
• • [3.1 .ckpt (Checkpoint)：不止是模型，更是"训练现场"的快照](#3.1 .ckpt (Checkpoint)：不止是模型，更是“训练现场”的快照)
  • [3.2 .pth (PyTorch)：一个"约定俗成"的旧爱](#3.2 .pth (PyTorch)：一个“约定俗成”的旧爱)
• [一张图看懂 三大主流格式的核心对比](#一张图看懂 三大主流格式的核心对比)
• [Hugging Face Hub 的"自动转换"：为什么你下载的模型越来越多是.safetensors？](#Hugging Face Hub 的“自动转换”：为什么你下载的模型越来越多是.safetensors？)
• 总结与展望：成为一个"靠谱"的AI模型管理者

前言：训练的终点，是"存档"的开始

经过数小时的等待，看着Loss曲线缓缓下降，你终于训练好了一个表现出色的AI模型。此刻，这个模型的所有"智慧"------那亿万个经过优化的参数------都还静静地躺在你的内存（或显存）里。

一旦你关闭程序，或者电脑断电，这些耗费了巨大算力才学到的"知识"都将烟消云散。

模型持久化 (Model Persistence)，也就是我们常说的"保存模型"或"存档"，就是将这些内存中的"智慧"，写入到硬盘文件中，以便未来可以随时加载回来，进行推理、继续训练或分享给他人。

但是，"存档"这件看似简单的事，背后却隐藏着格式的选择、效率的考量，甚至严重的安全风险。今天，我们就来当一次"AI档案管理员"，学习如何正确地管理这些珍贵的"数字大脑"。

第一章：PyTorch的"传统艺能"------ torch.save 与 Pickle 协议

介绍PyTorch原生的两种保存方式，并揭示其背后pickle协议的巨大安全隐患。

1.1 两种存档策略：只存"灵魂" vs. "灵魂+躯体"一起存

只保存模型权重 (State Dict)：

比喻：只保存一个人的"记忆和灵魂"（state_dict），不保存他的"身体结构"。

优点：文件小，灵活，安全。这是官方推荐的最佳实践。

缺点：加载时，你必须先拥有一个一模一样的"身体"（即模型的类定义），才能把"灵魂"加载进去。

保存整个模型：

比喻："灵魂"和"身体结构"一起打包保存。

优点：方便，加载时不需要模型的类定义。

缺点：文件大，耦合性强，存在严重的安全风险。

1.2 用代码实现两种存档与读档方式

本节概括：我们将定义一个极简的神经网络，然后分别用"只存权重"和"存完整模型"两种方式进行保存和加载，让你直观地看到它们的区别。

代码实现

import torch
import torch.nn as nn
import os

# --- 0. 准备工作：创建一个文件夹和定义一个简单模型 ---
os.makedirs('models', exist_ok=True)

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear1 = nn.Linear(10, 20)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(20, 5)

    def forward(self, x):
        return self.linear2(self.relu(self.linear1(x)))

# 实例化一个模型
model = SimpleModel()
print("--- 原始模型 (随机初始化权重) ---")
# 打印一个权重，看看它的初始值
print(model.linear1.weight[0, :5])


# ==========================================================
# 策略一：只保存"灵魂"(State Dict) - 官方推荐
# ==========================================================
print("\n--- 策略一：只保存权重 (State Dict) ---")

# 1. 保存
# model.state_dict() 会返回一个包含模型所有权重的字典
STATE_DICT_PATH = 'models/simple_model_weights.pt'
torch.save(model.state_dict(), STATE_DICT_PATH)
print(f"✅ 模型的权重已保存到: {STATE_DICT_PATH}")

# 2. 加载
# 首先，你需要重新创建一个一模一样的模型"躯体"
new_model_from_weights = SimpleModel()
# 然后，将保存的权重"灵魂"加载进去
new_model_from_weights.load_state_dict(torch.load(STATE_DICT_PATH))
# 设置为评估模式
new_model_from_weights.eval() 
print("✅ 权重已成功加载到新的模型实例中！")
print("--- 从权重加载的模型 ---")
print(new_model_from_weights.linear1.weight[0, :5])


# ==========================================================
# 策略二：保存"灵魂 + 躯体"(完整模型) - 不推荐
# ==========================================================
print("\n--- 策略二：保存完整模型 ---")

# 1. 保存
# 直接保存整个model对象
FULL_MODEL_PATH = 'models/simple_model_full.pt'
torch.save(model, FULL_MODEL_PATH)
print(f"✅ 完整的模型已保存到: {FULL_MODEL_PATH}")

# 2. 加载
# 直接加载文件，无需预先定义模型类
# (但前提是，你当前的代码环境能找到SimpleModel这个类的定义)
new_model_from_full = torch.load(FULL_MODEL_PATH)
new_model_from_full.eval()
print("✅ 完整的模型已成功加载！")
print("--- 从完整模型文件加载的模型 ---")
print(new_model_from_full.linear1.weight[0, :5])

print("\n🔍 对比发现，三种模型的权重完全一致，证明加载成功！")

【代码解读】

运行这段代码，你会清晰地看到：

策略一更"安全"，因为它要求你必须拥有SimpleModel的源代码才能恢复模型。你清楚地知道自己正在加载什么。

策略二看起来更"方便"，torch.load()一步到位。但它的代价是巨大的安全风险和较差的可移植性。

最佳实践：永远优先使用策略一（保存和加载state_dict）。

1.3 【高能预警】Pickle的"阿喀琉斯之踵"：恐怖的"代码注入"风险

torch.save的底层，使用的是Python的pickle模块。pickle的强大之处在于，它不仅能序列化数据，还能序列化代码对象。

这也带来了它最致命的弱点：一个恶意的.pt或.pth文件，可以在你加载它（torch.load）的时候，执行任意的、有害的计算机代码！

比如，一个黑客可以构造一个模型文件，你一加载，它就在后台悄悄地删除你的文件、或者上传你的私钥。

结论：永远不要加载来自你不信任来源的.pt或.pth文件！

第二章：新王登基 ------ safetensors 的崛起

介绍Hugging Face推出的safetensors格式，解释它为什么更安全、更快速，以及如何使用它。

为了解决Pickle的根本性安全问题，Hugging Face社区开发并力推一种新的模型序列化格式：.safetensors。

2.1 为什么需要safetensors？------ 安全、快速、跨平台

安全 (Safe)：这是它最重要的特性。

快速 (Fast)：它的加载速度，尤其是当模型只加载一部分（分片加载）时，远快于pickle。

零拷贝 (Zero-copy)：加载时可以直接将硬盘上的数据映射到内存，无需额外的反序列化开销。

2.2 safetensors 的核心思想：数据与代码的"彻底分离"

safetensors的实现原理非常简单：

文件结构：一个.safetensors文件，只包含两部分：

一个JSON头部：用纯文本描述了所有张量的元信息（名字、形状、数据类型、在文件中的位置）。

一个二进制数据块：包含了所有Tensor的、紧凑排列的、原始的二进制字节数据。

加载过程：加载时，程序首先安全地解析JSON头部，了解模型结构，然后直接从文件中的指定位置，将二进制数据块加载到内存中，构建成Tensor。

整个过程中，没有任何执行任意代码的机会。 它只处理数据，不处理代码。

2.3 用Hugging Face safetensors库进行安全的模型存取

我们将使用与上一个案例相同的简单模型，但这次，我们用safetensors来保存和加载它的state_dict，体验其安全、高效的特性。

代码实现

import torch
import torch.nn as nn
import os
# 首先，需要安装safetensors库
# pip install safetensors
from safetensors.torch import save_file, load_file

# --- 0. 准备工作：复用我们之前的SimpleModel ---
os.makedirs('models', exist_ok=True)

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear1 = nn.Linear(10, 20)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(20, 5)

    def forward(self, x):
        return self.linear2(self.relu(self.linear1(x)))

# 实例化一个模型
model_to_save_safely = SimpleModel()
print("--- 待保存的模型权重 (部分) ---")
print(model_to_save_safely.linear1.weight[0, :5])


# --- 1. 使用safetensors进行保存 ---
print("\n--- 使用 safetensors 保存 ---")

# safetensors只操作state_dict，符合最佳实践
weights = model_to_save_safely.state_dict()
SAFETENSORS_PATH = 'models/simple_model.safetensors'

# 调用 save_file 函数
save_file(weights, SAFETENSORS_PATH)
print(f"✅ 模型的权重已安全地保存到: {SAFETENSORS_PATH}")


# --- 2. 使用safetensors进行加载 ---
print("\n--- 使用 safetensors 加载 ---")

# 同样，先创建一个新的模型"躯体"
new_model_safe_load = SimpleModel()

# 调用 load_file 函数，它会返回一个state_dict
loaded_weights = load_file(SAFETENSORS_PATH)

# 将加载的权重应用到模型上
new_model_safe_load.load_state_dict(loaded_weights)
new_model_safe_load.eval()
print("✅ 权重已从.safetensors文件安全加载！")
print("--- 从safetensors加载的模型权重 (部分) ---")
print(new_model_safe_load.linear1.weight[0, :5])

print("\n🔍 对比发现，权重完全一致，我们完成了一次安全的模型存取！")

代码解读与优势分析】

你会发现，使用safetensors的流程和PyTorch官方推荐的"只存权重"策略几乎完全一样，学习成本极低。但它带来了根本性的改变：

绝对安全：load_file这个函数在设计上就不可能执行任何恶意代码。你可以放心地加载任何来源的.safetensors文件。

跨框架兼容：用PyTorch保存的.safetensors文件，可以被TensorFlow、JAX等其他框架直接加载，反之亦然。它是一种通用的"Tensor档案袋"。

加载飞快：对于巨大的模型文件（几十上百GB），safetensors的加载速度优势会变得极其明显，因为它避免了pickle缓慢的反序列化过程。

结论：在你的所有项目中，请养成使用safetensors.torch.save_file来替代torch.save保存权重的习惯。这是一个能让你受益终身的、专业且安全的工程实践。

第三章：江湖别传 ------ .ckpt 与 .pth 的故事

解释另外两种常见的文件后缀，让你对模型文件的"江湖"有一个更完整的认识。

3.1 .ckpt (Checkpoint)：不止是模型，更是"训练现场"的快照

.ckpt或.checkpoint文件，通常不仅仅包含模型的权重。它是一个**"检查点"**文件，除了model_state_dict，往往还保存了：

optimizer_state_dict：优化器（如Adam）的状态，包含了动量等信息，以便能从断点处恢复训练。

epoch：当前训练到了第几个周期。

loss：当前的损失值等。

一句话：.ckpt是用于"恢复训练"的，而.pt或.safetensors通常是用于"推理部署"的最终产物。

3.2 .pth (PyTorch)：一个"约定俗成"的旧爱

.pth是PyTorch早期教程中常用的文件后缀，它和.pt在技术上没有任何区别，都是使用pickle序列化的结果。只是一个社区的命名习惯问题。现在，.pt更为常见

一张图看懂 三大主流格式的核心对比

特性 .pt / .pth (Pickle) .safetensors .ckpt (Checkpoint)

核心技术 Python pickle JSON + Raw Binary Python pickle

安全性 🚨 极不安全 ✅ 非常安全 🚨 极不安全

加载速度 普通 ⚡️ 极快 (尤其是分片) 普通

主要用途 推理部署 (不推荐) 推理部署 (推荐) 恢复训练

内容物 权重 或 完整模型 通常只有权重 权重、优化器状态、训练进度等

Hugging Face Hub 的"自动转换"：为什么你下载的模型越来越多是.safetensors？

你可能已经注意到，现在在Hugging Face Hub上下载模型时，即使作者只上传了.pt文件，Hub通常也会自动为你提供一个.safetensors的下载选项。

这是Hugging Face为了社区安全，在后台做的一项重要工作。它会自动加载用户上传的pickle文件，提取出其中的权重，然后用.safetensors格式重新保存一份。当你使用from_pretrained方法时，如果检测到有.safetensors文件，它会优先加载这个安全版本。

总结与展望：成为一个"靠谱"的AI模型管理者

恭喜你！今天你从一个模型的使用者，成长为了一名懂得如何安全、高效地管理和分发模型的"档案管理员"。

✨ 本章惊喜概括 ✨

你掌握了什么？	对应的技能/工具
理解了两种存档策略	✅ 只存权重 vs. 存完整模型
洞悉了Pickle的风险	✅ 警惕.pt文件的"代码注入"攻击
掌握了未来的标准	✅ safetensors的安全与高效
分清了不同的"黑话"	✅ .ckpt用于恢复训练，.pth是旧称
在AI的开源世界里，模型的分享与协作是基石。而安全，则是这一切的底线。养成优先使用和分享.safetensors格式的习惯，不仅是对你自己负责，也是对整个社区负责。