AutoModel 是 Hugging Face Transformers 里我最常用的类没有之一。不管是加载别人训练好的开源模型还是把本地 checkpoint 塞进 pipeline一句AutoModel.from_pretrained()就搞定了。但等你自己开始训练自定义模型、想用同一套加载逻辑把模型存下来再读回去时多半会撞上一堆报错Unrecognized model type、Could not find modeling file、size mismatch每一行都让人头大。这篇文章就围绕“AutoModel 怎么保存与加载自定义模型”这件事把机制拆开讲清楚然后给你一套能直接落地复制的方案。适合这类场景的读者你在某个内部项目里改过注意力结构或者给模型加了自定义 head又或者只是想把一个非官方结构的模型用统一方式加载。我的目标是你看完之后能自己写一个模型文件、注册到 AutoModel、保存成 checkpoint再顺利读回来不靠网上零碎的帖子拼凑。1. 先搞清楚 AutoModel 到底是怎么“认人”的1.1 AutoModel 不是魔法它只是一个“路由表”很多人在自定义模型加载失败时第一反应是去改 Transformers 源码加一个 if 分支。这个方向一开始就偏了。AutoModel的本质不是模型实现而是一个工厂它通过配置文件里的字段找到应该实例化哪个类然后调用那个类的from_pretrained方法完成加载。具体来说AutoModel.from_pretrained()会先读取目标目录下的config.json从里面拿到model_type或architectures字段再到注册表里查对应的模型类。你自定义的模型如果不在这个注册表里它自然不认识。这就像打车软件派单乘客输入的终点是model_type司机池子是注册表乘客要去的具体地址才是你的模型类。你只告诉它目的地它能派单的前提是系统里得有这个司机。所以要让 AutoModel 认出自定义模型本质上只有两条路要么把你的模型类注册进它的路由表要么在config.json里给出一份“导航地图”auto_map让它按图索骥找到你的代码。两条路不冲突可以搭配使用。1.2 保存自定义模型时需要哪些“身家证明”既然 AutoModel 是靠 config 找类那么“保存模型”就远不止把state_dict存下来这么简单。一个能被 AutoModel 重新加载的自定义模型目录里至少要有这几样东西config.json模型的元信息文件也是 AutoModel 的导航入口。pytorch_model.bin或model.safetensors权重文件。modeling_my_model.py模型类的定义代码AutoModel 需要执行它才能拿到你的类。可能还有tokenizer、preprocessor_config.json等取决于你的任务类型。很多人只保存了权重然后拿一个全新的模型类去load_state_dict也能跑但代价是加载逻辑散落在代码里每个使用方都要知道“模型类是什么、参数怎么传”。这不叫“能用 AutoModel 保存加载”叫“碰巧能跑”。如果你想让团队里任何一个人拿到这个目录就能训练、推理、部署那么 config、代码、权重三者必须齐全并且路径关系稳定。2. 动手之前先搭好自定义模型的目录与配置2.1 一种简单可靠的目录结构先说结论我建议你在项目里用下面这种结构放自定义模型相关的文件my_project/ ├── modeling_my_model.py ├── configuration_my_model.py └── save_and_load_demo.py为什么把配置类和模型类拆开因为config和model的职责不同config负责记录超参和元信息model负责定义网络结构。如果你把它们写在同一个文件里不是不行但加载时auto_map的路径写法会稍微绕一点。拆开之后职责清晰也方便你后续把模型推到某个模型仓库时按社区惯例管理。实际保存时生成的目标目录结构大概是这样my_custom_model/ ├── config.json ├── modeling_my_model.py ├── configuration_my_model.py ├── model.safetensors └── vocab.json这里的modeling_my_model.py和configuration_my_model.py是你在项目里写的原始文件save_pretrained会帮你把它们复制到输出目录至少它会为 remote code 模式做这一步。2.2 configuration_my_model.py 里必须写清楚什么你的配置类应该继承PretrainedConfig并且要显式设置model_type。这个字段的值就是 AutoModel 以后识别你的模型的钥匙。比如from transformers import PretrainedConfig class MyModelConfig(PretrainedConfig): model_type my_model def __init__( self, vocab_size32000, hidden_size768, num_hidden_layers12, num_attention_heads12, intermediate_size3072, max_position_embeddings1024, initializer_range0.02, **kwargs, ): super().__init__(**kwargs) self.vocab_size vocab_size self.hidden_size hidden_size self.num_hidden_layers num_hidden_layers self.num_attention_heads num_attention_heads self.intermediate_size intermediate_size self.max_position_embeddings max_position_embeddings self.initializer_range initializer_range这里有一个很多新手会漏掉的点model_type一定要在__init__外定义或者说至少保证它作为类属性存在。因为save_pretrained在写出config.json时会读取这个属性。如果你的model_type只通过kwargs传入而不在类上声明序列化时可能丢掉它加载的时候就抓瞎。另外kwargs里一定要调用super().__init__(**kwargs)。很多PretrainedConfig的内置字段torch_dtype、transformers_version、_name_or_path等都需要经过父类初始化才能正确落盘。漏掉这一行from_pretrained可能静默丢失某些元信息后排查起来很麻烦。2.3 modeling 文件里必须包含的三个关键接口自定义模型文件的核心是你自己的nn.Module实现但要能被 AutoModel 用起来还缺三样东西模型类本身通常命名如MyModel。必须在文件顶部from transformers import PreTrainedModel并继承PreTrainedModel而不是nn.Module。必须在类里定义config_class MyModelConfig并在__init__里调用super().__init__(config)。原因不复杂PreTrainedModel是 Transformers 给所有模型提供的基础设施save_pretrained、from_pretrained、push_to_hub、generate这些方法全都依赖它。你如果直接继承nn.Module那么AutoModel.from_pretrained即使找到了你的类调用的也是父类的通用加载逻辑很多预处理和后处理逻辑对不上迟早出问题。模型的前向函数返回值建议保持 Transformers 的惯例返回一个包含loss和logits的对象或者直接返回ModelOutput。这不是强制要求但如果你后面接 Trainer、pipeline或者想统一用AutoModelForSequenceClassification之类的类做封装返回值结构一致可以省掉一堆兼容代码。3. 保存与导出的完整流程3.1 用 save_pretrained 保存模型的完整动作假设你已经训练好了一个自定义模型model对应的配置是config。保存时不要只调用torch.save(model.state_dict(), ...)那会把 config 和代码文件丢掉。正确的做法是model.save_pretrained(./my_custom_model) model.config.save_pretrained(./my_custom_model)如果你把config绑定到了模型上也就是model.config config第一个调用就会自动把 config 写进去第二个调用可以省略。但显式写下config.save_pretrained有个额外好处你可以临时修改 config 里的某些字段再保存比如把max_position_embeddings调低或者把architectures字段改成你想要的名字。safetensors是另一个建议。从某个版本开始save_pretrained默认会尝试保存为model.safetensors。如果你更习惯用老的 bin 格式也可以显式传入safe_serializationFalse。我个人的建议是多设备推理、生产环境优先用 safetensors加载更快也避免 pickle 反序列化带来的安全隐患。你的加载代码并不需要关心它是 bin 还是 safetensorsAutoModel 会自动识别。保存完之后建议打开config.json检查一眼。重点看这几个字段在不在model_type值等于你在 config 类里定义的字符串。architectures值是模型类的名字数组比如[MyModel]。transformers_version用于追踪版本兼容性。这些字段里model_type是 AutoModel 认路的关键architectures是给库内注册表用的。两个都别丢。3.2 把自定义模型注册进 AutoModel 的两种方式方式一使用auto_map远程代码。这种方式需要在config.json里显式声明让 AutoModel 知道去哪个文件、找哪个类。下面是一个示例配置{ model_type: my_model, auto_map: { AutoConfig: configuration_my_model.MyModelConfig, AutoModel: modeling_my_model.MyModel } }然后加载时你必须传trust_remote_codeTrue并指定config所在的目录from transformers import AutoConfig, AutoModel config AutoConfig.from_pretrained(./my_custom_model, trust_remote_codeTrue) model AutoModel.from_pretrained(./my_custom_model, configconfig, trust_remote_codeTrue)注意auto_map里配置的路径是相对于加载路径的模块路径不是 Python 包路径。也就是说如果config.json在my_custom_model目录下它会去my_custom_model/modeling_my_model.py和my_custom_model/configuration_my_model.py里找对应类而不是去 site-packages 里找。这很关键很多人在这踩坑。方式二直接使用register类方法。如果你只是在自己项目里用不想每次加载都传trust_remote_code可以直接注册from transformers import AutoConfig, AutoModel AutoConfig.register(my_model, MyModelConfig) AutoModel.register(MyModelConfig, MyModel)注册之后AutoModel 的路由表里就有你了。这时加载模型就不需要trust_remote_code也不需要auto_map只要目录里有config.json和权重文件就行。缺点是这个注册动作必须在加载前执行过其他使用者也得在你的环境里先跑一遍注册代码。两种方式怎么选我的判断是如果你要把模型发布给其他人用或者放到模型仓库用auto_map更规范。如果你只是在本地训练和推理循环里复用用register更简单少了一层文件映射的复杂度。多数内部项目register已经足够了。3.3 加载时的参数到底该怎么传AutoModel.from_pretrained里最常被忽视的两个参数是trust_remote_code和strict。trust_remote_code的含义是允许 AutoModel 加载并执行自定义代码文件。这就像一个“可执行文件”许可你确认目录里的.py文件是可信的才会去执行。本地加载自己写的代码问题不大。但如果你下载了一个陌生人的模型并且它的 config 里有auto_map执行它之前务必检查代码内容。这是社区里反复强调的安全意识但实践里真的有人不看。strict参数决定load_state_dict是否要求权重 key 完全匹配。默认True也就是模型代码定义的层和权重文件里的 key 必须一一对应。自定义模型改结构后加载老权重或者加载了和 config 不一致的 checkpoint最常见的报错就是size mismatch和unexpected key。如果你确实想部分加载可以传strictFalse但要明确这可能导致你模型里有随机初始化的层语义上是“加载了但没完全加载”。一个值得养成的习惯是加载后做一次输出验证用同一个输入跑一次旧模型和新加载的模型对比输出张量是否一致。很多静默错误比如 config 里的hidden_size写错了导致权重被广播式复用都是靠这种 sanity check 暴露的。4. 实操示例从自定义实现到 AutoModel 完整加载4.1 示例模型设计下面是我在一个模拟项目里用过的例子结构很简单一个 embedding 层、一个 Transformer encoder 层、一个分类头。为了让这个例子有区分度我没用现成的BertModel直接拼而是写了一个自定义初始化方式的模块。# modeling_my_model.py import torch import torch.nn as nn from transformers import PreTrainedModel from configuration_my_model import MyModelConfig class MyModel(PreTrainedModel): config_class MyModelConfig def __init__(self, config): super().__init__(config) self.embedding nn.Embedding(config.vocab_size, config.hidden_size) self.encoder nn.TransformerEncoder( nn.TransformerEncoderLayer( d_modelconfig.hidden_size, nheadconfig.num_attention_heads, dim_feedforwardconfig.intermediate_size, batch_firstTrue, ), num_layersconfig.num_hidden_layers, ) self.classifier nn.Linear(config.hidden_size, 2) self.post_init() def forward(self, input_ids, attention_maskNone, labelsNone): hidden self.embedding(input_ids) hidden self.encoder(hidden) pooled hidden.mean(dim1) logits self.classifier(pooled) loss None if labels is not None: loss nn.functional.cross_entropy(logits, labels) return {loss: loss, logits: logits}注意这里调用了post_init()这个方法会在初始化后做权重初始化并绑定config。如果你需要新的随机初始化策略也应该在post_init前设置好。4.2 训练并保存 checkpoint接下来是保存流程。为了体现“完整”我把 tokenizer、config、模型三个部分拆开来演示from transformers import AutoTokenizer from configuration_my_model import MyModelConfig from modeling_my_model import MyModel config MyModelConfig( vocab_size1000, hidden_size128, num_hidden_layers2, num_attention_heads4, intermediate_size512, max_position_embeddings64, ) model MyModel(config) tokenizer AutoTokenizer.from_pretrained(bert-base-uncased) output_dir ./my_custom_model model.save_pretrained(output_dir, safe_serializationTrue) model.config.save_pretrained(output_dir) tokenizer.save_pretrained(output_dir).save_pretrained之后output_dir的config.json里应该能看到model_typemy_model。如果你的模型类没有自动带出这个字段或者你想让其他人在本地不执行注册代码也能加载就在config.json里手动补一份auto_map。注意save_pretrained本身不会帮你生成auto_map这是需要你主动加的。4.3 用 AutoModel 加载并做一致性验证接下来的加载演示分两个层级。先是最朴素的目录里没有auto_map但代码已经在当前 Python 进程里注册好了。from transformers import AutoConfig, AutoModel AutoConfig.register(my_model, MyModelConfig) AutoModel.register(MyModelConfig, MyModel) config AutoConfig.from_pretrained(./my_custom_model) model AutoModel.from_pretrained(./my_custom_model, configconfig) model.eval() with torch.no_grad(): logits model(torch.tensor([[1, 2, 3]]))[logits] print(logits.shape) # torch.Size([1, 2])如果目录里有auto_map也简单config AutoConfig.from_pretrained(./my_custom_model, trust_remote_codeTrue) model AutoModel.from_pretrained(./my_custom_model, configconfig, trust_remote_codeTrue)两种方式加载之后模型输出应该完全一致。如果输出不一致说明 class 路径或者 config 序列化环节出了问题。排查时先看config.json里的model_type和代码里定义的是否一致再看architectures字段里的类名和代码里的类名是否一致。5. 常见问题与排查技巧实录5.1 典型报错速查表我做了一张小表按我自己的踩坑频率排序。不一定全面但覆盖了 80% 的自定义模型加载问题。报错特征常见原因解决办法Unrecognized model typemodel_type没有注册调用AutoConfig.register或在config.json里写auto_mapCould not find modeling fileauto_map路径配置错误确认路径是相对 checkpoint 目录的模块路径并确认.py文件已被复制到目录ImportError: cannot import name类名和auto_map里的名字不一致打开目标目录下的.py文件核对类名size mismatch for embedding.weightconfig 参数和权重维度不一致检查vocab_size、hidden_size等是否保存时与加载时一致Unexpected key(s) in state_dict模型结构和权重 key 对不上加载时传strictFalse确认缺失层或重新初始化对齐的结构加载成功但输出明显异常config 参数字段名错误导致多个层复用相同参数打印 config 和模型参数量逐个核对字段名trust_remote_codeTrue后仍然报错当前 Transformers 版本过旧不支持该参数升级 Transformers或改用register方式OSError: model.safetensors is not a valid file手动改名导致的文件格式不匹配用save_pretrained(safe_serializationTrue)重新保存5.2 一条能快速定位问题的排查路径我以前遇到问题习惯从最外层往里走。现在推荐的排查路径是先看config.json是否存在且内容完整。缺了这个文件AutoModel 连路由都没法做。再看model_type是否等于你 config 类里定义的字符串。这一个字段错了后面所有的报错都是白看。看目录里有没有对应的.py文件。如果你用了auto_mapsave_pretrained不会自动把所有 py 文件都复制过去你要么手动复制要么保存时指定model_file_name和config_file_name。看权重文件是否存在并检查它和 config 的参数维度是否匹配。最简单的方式是直接用torch.load读一下观察 key 和 shape。最后一步再去看AutoModel的报错堆栈。多数情况下到这一步会发现前面某个环节漏了。5.3 集成到训练框架时的额外注意点很多人不是单独跑上面的 demo而是把自定义模型塞进Trainer或自己的训练循环里。这里有几个坑我逐一说明。第一Trainer在保存 checkpoint 时会调用model.save_pretrained(save_directory)。如果你的模型文件里没有带上configuration_my_model.py和modeling_my_model.py那么拿这个 checkpoint 给别人加载时就会有挫败感。解决办法是在保存前把这两个.py文件复制到输出目录。或者更优雅一点重写save_pretrained让它在保存 config 和权重之前先把源文件复制过去。第二多卡device_map加载时AutoModel.from_pretrained的device_mapauto会触发模型拆分。自定义模型如果内部结构不是标准的nn.Module嵌套或者某些层用了不可分片的实现拆分时可能报错。这种情况我要么退回device_mapNone单卡加载要么把自定义算子改成可分片的形式。第三如果你自定义模型里包含随机 dropout、或者非确定性采样层跨版本加载后重新训练指标可能对不上。保存时最好把transformers_version字段留下来加载时也确认代码文件版本和保存时一致。这个字段PretrainedConfig会自动记录但如果你手动构造 config 并且不调用save_pretrained就可能缺失。5.4 我踩过的几次坑第一次踩坑是我把config_class写错位置放在了文件底部而不是模型类里。加载时AutoModel.from_pretrained找到了类但拿不到 config 类于是退回到默认的PretrainedConfig一堆参数全部丢失。这种问题不会立刻报错而是等到模型前向计算时才发现维度全乱了排查起来格外痛苦。现在我的习惯是写完modeling文件后先单独跑一个“保存再加载”的最小测试确认 config 字段一致再继续往下写训练逻辑。第二次踩坑是auto_map路径写错。我把路径写成了my_custom_model/modeling_my_model.MyModel但实际config.json就在my_custom_model目录下它加载时会把目录名也拼接进去结果自然找不到。后来我统一约定auto_map里只写相对于 checkpoint 目录的模块名和类名不写目录前缀。这个约定沿用至今没再出过类似问题。第三次踩坑比较隐蔽模型里用了自定义的激活函数但这个函数写在训练脚本里没放进modeling_my_model.py。结果保存时一切正常一旦换一台机器用AutoModel.from_pretrained加载直接报NameError。从那以后我要求所有模型相关代码包括自定义初始化、激活函数、损失函数全部放回模型文件里保持目录自包含。最后分享一点个人体会做了几年的模型训练和部署我最大的感受是加载逻辑这件事越早统一越好。训练时顺手把 config、模型文件、tokenizer 一起保存比事后补档案省心太多。AutoModel 这套设计虽然初看有些绕但它的价值在于“约定优于配置”只要你的文件结构和字段命名符合它的约定它就能像加载官方模型一样加载你的自定义模型这样团队协作、模型上线、后续迭代都会顺很多。如果你也在做类似的事情建议从最小用例开始先把保存加载闭环跑通再逐渐加复杂度这套体系会给你带来稳定的复现体验。
