Transformers save_pretrained 怎样保存可迁移模型

要让 Transformers 模型可以迁移到另一台机器，核心不是只复制一个权重文件，而是把模型对象和 tokenizer 对象保存到同一个目录，再在新环境中分别恢复。模型侧负责权重与 config.json，tokenizer 侧负责词表、特殊 token 和分词配置。

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

一、先把迁移对象拆成模型与 tokenizer

模型和 tokenizer 是两个有边界的对象。模型需要知道网络结构、隐藏层配置和权重文件；tokenizer 则决定文本怎样变成 input_ids，还包含特殊 token、词表及其映射。只保存模型可能得到一个能加载的网络，却无法保证新环境生成相同的输入。

二、用同一个输出目录保存完整组合

推荐把训练完成的模型和实际参与训练的 tokenizer 指向同一个输出目录。下面的代码展示保存关系；注释只标出关键职责，不代表已经在本机执行。

from pathlib import Path
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_id = "distilbert/distilbert-base-uncased"
output_dir = Path("artifacts/sentiment-model")

# 模型与 tokenizer 必须来自同一套输入约定
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2)

# 目录不存在时由 Transformers 创建；模型侧写权重和配置
model.save_pretrained(output_dir)
# tokenizer 侧写词表、特殊 token 与分词器配置
tokenizer.save_pretrained(output_dir)

保存后目录可能包含 config.json、model.safetensors（或分片文件）以及 tokenizer 的 JSON、词表和映射文件。具体文件名由模型架构、Transformers 版本和 tokenizer 类型决定，所以迁移时应整体复制目录，不要只挑一个看起来最大的文件。

三、在新进程按目录恢复并检查匹配关系

恢复时先加载 tokenizer，再加载模型，调用方就能沿用训练时的编码方式。from_pretrained 接受保存目录，并会读取其中的配置；它不是要求你重新拼装模型结构。

from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_dir = "artifacts/sentiment-model"

# 两个对象都从同一目录恢复，避免词表与权重来自不同版本
tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
model = AutoModelForSequenceClassification.from_pretrained(model_dir, local_files_only=True)

# 迁移前可用这个边界做快速人工核对
if len(tokenizer) != model.get_input_embeddings().num_embeddings:
    raise ValueError("tokenizer 词表大小与模型输入嵌入不匹配")

如果训练期间调用过 tokenizer.add_tokens()，还需要在保存前确认模型的输入嵌入已经通过对应方式扩展，否则 tokenizer 文件虽然更新，模型仍可能无法承接新增 token。这个检查比比较文件大小更有意义。

四、处理分片、覆盖与不该混放的文件

大模型保存时可以用 max_shard_size 控制单个权重分片的上限，例如 model.save_pretrained(output_dir, max_shard_size="2GB")。分片不是多个独立模型，必须和生成的索引文件一起迁移。

重复保存到同一目录前，先确认目录用途。最稳妥的做法是为每次训练使用独立目录，保存完成后再把说明文件、评估结果或推理脚本放在旁边；不要把临时缓存、另一模型的 tokenizer 或未经记录的手工词表混入模型目录。若改过 tokenizer 或分类头，最好在目录外记录基座模型、标签映射和 Transformers 版本，便于回滚和复现。

相关问题

只保存 model.state_dict() 可以迁移吗？

可以作为底层权重备份，但它不包含完整的 Transformers 配置和 tokenizer 状态。面向可复用目录时，优先使用 save_pretrained 组合。

tokenizer 一定要和模型放在同一个目录吗？

不一定，但同目录最容易保证版本对应。若分开保存，必须把两者的来源、版本和词表变更记录清楚，恢复时显式传入正确路径。