这里说的 TF 是 TensorFlow不是那张插在手机或单片机上的存储卡——先把这句放在最前头免得点进来的朋友误会。我做模型训练和线上部署这些年SavedModel 这个格式前前后后用过几百次从 TF 1.x 那套tf.saved_model.builder.SavedModelBuilder到 TF 2.x 的tf.saved_model.save再到 Keras 的model.save(filepath)被它的目录结构、版本戳、签名覆盖恶心到的次数两只手都数不过来。真正让人抓狂的不是第一次保存而是同一个训练流程里多次保存每个 epoch 存一次、最佳权重存一次、训练结束再导一次目录里就会莫名其妙多出assets、多出几组variables分片甚至出现加载时报 CRC 校验失败、签名找不到、自定义层认不出来这类抽风现象。这篇就把多次保存模型这条路上的坑从文件结构、保存 API 选择、目录管理到排查手段一次性讲透适合正在做模型训练、需要定期落盘、又要给别人交付可加载模型的同学。1. 先搞清楚 SavedModel 这个目录里究竟装了什么1.1 一个目录四类关键文件SavedModel 从来就不是一个单文件它是一个目录。你用tf.saved_model.save(model, path)之后得到的典型结构是这样my_model/ ├── saved_model.pb ├── variables/ │ ├── variables.data-00000-of-00001 │ └── variables.index ├── assets/ └── fingerprint.pbsaved_model.pb用 protobuf 序列化出来的 MetaGraphDef里面装的是计算图结构、签名定义、节点名称、输入输出张量的元信息。它不含权重。variables/真正的权重。variables.data-00000-of-00001存数值variables.index存每个变量名到分片位置的索引和 shape。assets/外部资源文件比如词表、查找表、自定义的映射文件。你如果用tf.lookup.StaticVocabularyTable之类的词表就落在这里。fingerprint.pbTF 2.7 之后引入的指纹文件用来校验导出的图与变量版本是否一致。这个拆分是 SavedModel 的核心设计图和权重分离。好处很实在——同一张图可以配不同权重权重也能单独做增量更新或热替换。但坏处同样明显只要你在多次保存中把这两部分里的任意一部分覆盖得不干净就会出现图和权重对不上号的诡异问题而且报错信息往往指向别处让人怀疑人生。注意fingerprint.pb在很多老版本 TF 里根本没有如果你用新版本保存、老版本加载官方不保证兼容。这也是多次保存场景里最容易忽略的一处版本陷阱。1.2 它和 checkpoint、H5 到底差在哪很多人纠结保存格式我把三者放在一起对比你自己看格式存什么适合场景多次保存的风险点Checkpoint只存变量权重训练中断续训、快速落盘不带图加载要重建结构H5结构权重优化器快速原型、Keras 内部用自定义对象多时难加载SavedModel图权重签名线上部署、跨语言推理目录残留、签名覆盖结论其实很清晰训练中途要频繁保存用 Checkpoint 最省事要交付给别人做推理部署才用 SavedModel。如果你把 SavedModel 当成训练时的每步落盘手段那基本上就是在给自己挖坑因为每次保存都要重新写一遍图和变量多次保存带来的目录膨胀和残留问题会被放大。1.3 签名这个东西到底管什么签名signatures是 SavedModel 的入口约定。你保存的时候没指定签名TF 会自动帮你加一个默认的serving_default把模型的可调用部分包进去。问题是如果你反复用tf.saved_model.save覆盖同一个目录而不同次传入的signatures字典 key 不一样后一次的签名会顶掉前一次旧签名直接消失。签名里保存的是 ConcreteFunction如果这个函数闭包捕获了额外的变量或常量整个图就会被拖进saved_model.pb模型体积会明显膨胀。所以多次保存时签名不是可有可无的装饰而是每次都该显式确认的东西。2. 多次保存模型最容易踩的四个坑2.1 坑一往同一个路径反复写旧文件残留在目录里这是最高频、也最隐蔽的坑。tf.saved_model.save或model.save(path)默认行为是在目标目录里覆盖写入但它并不保证把上一版遗留的多余文件删掉。举个我自己踩过的实例第一次保存时模型变量比较多TF 自动分成了 3 个数据分片variables/ ├── variables.data-00000-of-00003 ├── variables.data-00001-of-00003 └── variables.data-00002-of-00003第二次保存时模型变小了比如剪枝之后只生成 1 个分片variables/ ├── variables.data-00000-of-00001这时候variables/目录里就同时存在着新写的00000-of-00001和旧的00001-of-00003、00002-of-00003。加载时 TF 会去读 index 文件index 指向的却是旧的00001-of-00003一读就报数据校验失败或者干脆加载出一个错误形状的权重。规避方法有两个保存前先清空目标目录最直接import shutil, os def clean_save(model, path): if os.path.exists(path): shutil.rmtree(path) model.save(path)每次保存用独立的子目录把版本信息写进路径从根上避免覆盖import os def versioned_save(model, base_dir, tag): save_dir os.path.join(base_dir, tag) os.makedirs(save_dir, exist_okTrue) model.save(save_dir) return save_dir实测下来第二种更稳因为它天然保留了历史版本回溯也方便。2.2 坑二自定义层、自定义损失加载时报 unknown class这个坑在多次保存里特别典型你在 notebook 里定义了一个MyAttentionLayer训练时保存没问题加载时却报ValueError: Unknown layer: MyAttentionLayer。原因很直白——Keras 加载 SavedModel 时需要能重新实例化你的类而这个类只存在于你当前的 Python 进程里TF 自然找不到。解决有两条路import tensorflow as tf # 方式一注册装饰器最推荐 tf.keras.utils.register_keras_serializable(packagemy_pkg) class MyAttentionLayer(tf.keras.layers.Layer): def __init__(self, units64, **kwargs): super().__init__(**kwargs) self.units units def get_config(self): config super().get_config() config.update({units: self.units}) return config# 方式二加载时显式传 custom_objects model tf.keras.models.load_model( my_model, custom_objects{MyAttentionLayer: MyAttentionLayer} )关键点在于get_config必须把自定义的构造参数都吐出来否则重建时参数丢失加载出来的模型结构和原来的不一样权重对不上会直接报 shape 不匹配。2.3 坑三签名被后一次保存覆盖线上找不到入口我在一次灰度发布里遇到过新版本模型用tf.saved_model.save保存到与旧版相同的目录结果旧版依赖的predict签名没了线上服务报Signature predict does not exist。原因是每次保存时signatures的 key 集合不同后一次会整体替换。正确做法是每次保存都显式、完整地给出签名并且保持一致import tensorflow as tf class MyModule(tf.Module): def __init__(self, model): super().__init__() self.model model tf.function(input_signature[ tf.TensorSpec(shape[None, 224, 224, 3], dtypetf.float32, nameimages) ]) def predict(self, images): return {scores: self.model(images, trainingFalse)} module MyModule(model) tf.saved_model.save(module, save_path, signatures{predict: module.predict})用tf.Module或者tf.function(input_signature...)把签名固定下来既能避免覆盖也能让导出的图更干净不至于把训练分支也塞进去。2.4 坑四训练图和推理图混着存体积翻倍还容易加载错很多人为了方便直接保存model本身结果 Keras 会把训练相关的节点dropout、BN 的 training 分支、损失函数、优化器状态全都序列化进图。多次保存后saved_model.pb的体积可能从几十 MB 涨到几百 MB而且加载时如果用错了入口推理结果会带上训练时的随机性。建议在保存推理模型时明确走一个干净的子模型或者用tf.Module包一层把trainingFalse固定进去这样每次保存产出的图都是同构的多次保存也不会互相干扰。3. 反复保存模型的正确姿势3.1 先分清tf.saved_model.save和model.save这两个 API 名字像、行为不同我列个表帮你决策API保存内容适合多次保存建议model.save(path)Keras 格式可含 H5 或 SavedModel训练完交付用带版本子目录tf.saved_model.save(obj, path)纯 SavedModel需给签名部署、跨语言显式签名清目录tf.train.Checkpoint只存变量训练中途落盘天然支持多份最安全我个人的习惯是训练循环里只写 Checkpoint训练结束时统一导一次 SavedModel。这样既避免频繁重写图也把多次保存的问题压缩到一次。3.2 目录命名策略让版本自己说话多人协作或者长期项目里我强推目录带版本戳的做法import os, time, shutil def save_with_version(model, root, prefixmodel, keep_last3): tag time.strftime(%Y%m%d-%H%M%S) save_dir os.path.join(root, f{prefix}_{tag}) os.makedirs(save_dir, exist_okTrue) model.save(save_dir) # 只保留最近 N 个防止磁盘被撑爆 versions sorted([ d for d in os.listdir(root) if d.startswith(prefix) ]) for old in versions[:-keep_last]: shutil.rmtree(os.path.join(root, old), ignore_errorsTrue) return save_dir这套逻辑解决了三个问题多版本共存、覆盖残留、磁盘膨胀。实测在 30 多天的连续训练里磁盘占用始终稳定在几个 G 以内。提示保留数量keep_last不要设太大SavedModel 单份动辄上百 MB尤其带 embedding 的模型几百 MB 很常见。3.3 加载时怎么选对签名加载 SavedModel 后第一件事是打印签名确认入口对import tensorflow as tf loaded tf.saved_model.load(model_20240101-120000) print(list(loaded.signatures.keys())) # [predict] infer loaded.signatures[predict] # 看输入输出结构避免喂错 shape print(infer.structured_input_signature) print(infer.structured_outputs)如果签名里有多个入口比如predict和embedding务必确认线上调用的是哪一个不要靠猜。多次保存场景下不同版本签名集合不一致是常态加载前先列一遍签名能省掉很多排查时间。4. 常见问题排查速查表4.1 报错信息与对应病灶报错/现象最可能的原因处理方式Data loss: not an sstable旧分片残留index 指向错文件清空目录重存Unknown layer: XXX自定义类未注册加注册装饰器或 custom_objectsSignature xxx does not exist签名被后续保存覆盖每次保存显式给同一组签名模型体积异常大训练分支被打进图用 tf.Module 包推理入口加载后推理结果随机BN/Dropout 走 training 分支固定 trainingFalse加载报 shape 不匹配结构变了但权重是老版本检查 get_config 完整返回4.2 三步定位法遇到 saved-model 加载异常我一般按下面三步走命中率很高看目录ls -R model_dir重点看variables/里分片数量是否和variables.index记录的一致多余的旧分片是重灾区。看签名加载后立刻list(loaded.signatures.keys())确认入口存在、输入 shape 和线上一致。看加载日志把 TF 日志等级调到 INFO观察它读到了哪个saved_model.pb、哪个变量文件很多问题在日志里直接写明了。注意排查时不要直接在线上目录上操作先把模型目录整个拷到本地一份避免误删线上文件。5. 一次完整的多次保存实操复盘假设我们要训练 10 个 epoch要求每个 epoch 落一次 Checkpoint第 5 个 epoch 和最后一个 epoch 各导出一次可部署的 SavedModel同时保留最近 3 个版本。完整流程如下import os import shutil import tensorflow as tf ROOT ./runs CKPT_DIR os.path.join(ROOT, ckpt) SM_DIR os.path.join(ROOT, saved_model) os.makedirs(CKPT_DIR, exist_okTrue) os.makedirs(SM_DIR, exist_okTrue) model tf.keras.Sequential([ tf.keras.layers.Input(shape(28, 28, 1)), tf.keras.layers.Conv2D(32, 3, activationrelu), tf.keras.layers.Flatten(), tf.keras.layers.Dense(10, activationsoftmax), ]) optimizer tf.keras.optimizers.Adam() ckpt tf.train.Checkpoint(modelmodel, optimizeroptimizer) ckpt_manager tf.train.CheckpointManager( ckpt, CKPT_DIR, max_to_keep3 ) def export_saved_model(model, base_dir, tag): 每次用独立子目录导出避免覆盖残留 save_dir os.path.join(base_dir, tag) if os.path.exists(save_dir): shutil.rmtree(save_dir) os.makedirs(save_dir, exist_okTrue) class ServingModule(tf.Module): def __init__(self, inner): super().__init__() self.inner inner tf.function(input_signature[ tf.TensorSpec(shape[None, 28, 28, 1], dtypetf.float32, nameimages) ]) def predict(self, images): return {probs: self.inner(images, trainingFalse)} module ServingModule(model) tf.saved_model.save( module, save_dir, signatures{predict: module.predict} ) return save_dir # 训练循环 for epoch in range(10): # 这里用假数据演示真实场景替换成你的 dataset x tf.random.normal([32, 28, 28, 1]) y tf.random.uniform([32], maxval10, dtypetf.int32) with tf.GradientTape() as tape: logits model(x, trainingTrue) loss tf.keras.losses.sparse_categorical_crossentropy(y, logits) loss tf.reduce_mean(loss) grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) ckpt_manager.save() # 每轮落一次 Checkpoint print(fepoch {epoch}, loss{loss.numpy():.4f}) if epoch 4 or epoch 9: path export_saved_model( model, SM_DIR, fv{epoch:02d} ) print(fexported - {path}) # 验证加载 loaded tf.saved_model.load(os.path.join(SM_DIR, v09)) infer loaded.signatures[predict] out infer(tf.random.normal([1, 28, 28, 1])) print(probs shape:, out[probs].shape)这段代码里有几个刻意设计的地方值得说一说Checkpoint 管训练落盘SavedModel 只管部署导出两者职责不混。多次保存的问题被彻底隔离到导出环节。export_saved_model里先rmtree再建目录杜绝了旧分片残留。这是我踩坑之后固定下来的写法。签名用tf.function(input_signature...)显式声明trainingFalse写死在函数体里保证每次导出的图同构线上加载结果稳定。用版本号做子目录多版本共存、可回溯加载时选版本目录即可。我个人在长期项目里的体会是SavedModel 本身设计没问题问题都出在人身上——图省事往一个目录反复写、不给签名、不清理历史文件。把导出环节当成一次正式的发布每次清目录、显式签名、独立版本目录这些坑基本就跟你无缘了。如果你现在正被not an sstable或者Unknown layer折磨先别急着改代码去variables/目录里数一数分片数量和文件名后缀十有八九能找到答案。
TensorFlow SavedModel 多次保存模型避坑与工程实践
这里说的 TF 是 TensorFlow不是那张插在手机或单片机上的存储卡——先把这句放在最前头免得点进来的朋友误会。我做模型训练和线上部署这些年SavedModel 这个格式前前后后用过几百次从 TF 1.x 那套tf.saved_model.builder.SavedModelBuilder到 TF 2.x 的tf.saved_model.save再到 Keras 的model.save(filepath)被它的目录结构、版本戳、签名覆盖恶心到的次数两只手都数不过来。真正让人抓狂的不是第一次保存而是同一个训练流程里多次保存每个 epoch 存一次、最佳权重存一次、训练结束再导一次目录里就会莫名其妙多出assets、多出几组variables分片甚至出现加载时报 CRC 校验失败、签名找不到、自定义层认不出来这类抽风现象。这篇就把多次保存模型这条路上的坑从文件结构、保存 API 选择、目录管理到排查手段一次性讲透适合正在做模型训练、需要定期落盘、又要给别人交付可加载模型的同学。1. 先搞清楚 SavedModel 这个目录里究竟装了什么1.1 一个目录四类关键文件SavedModel 从来就不是一个单文件它是一个目录。你用tf.saved_model.save(model, path)之后得到的典型结构是这样my_model/ ├── saved_model.pb ├── variables/ │ ├── variables.data-00000-of-00001 │ └── variables.index ├── assets/ └── fingerprint.pbsaved_model.pb用 protobuf 序列化出来的 MetaGraphDef里面装的是计算图结构、签名定义、节点名称、输入输出张量的元信息。它不含权重。variables/真正的权重。variables.data-00000-of-00001存数值variables.index存每个变量名到分片位置的索引和 shape。assets/外部资源文件比如词表、查找表、自定义的映射文件。你如果用tf.lookup.StaticVocabularyTable之类的词表就落在这里。fingerprint.pbTF 2.7 之后引入的指纹文件用来校验导出的图与变量版本是否一致。这个拆分是 SavedModel 的核心设计图和权重分离。好处很实在——同一张图可以配不同权重权重也能单独做增量更新或热替换。但坏处同样明显只要你在多次保存中把这两部分里的任意一部分覆盖得不干净就会出现图和权重对不上号的诡异问题而且报错信息往往指向别处让人怀疑人生。注意fingerprint.pb在很多老版本 TF 里根本没有如果你用新版本保存、老版本加载官方不保证兼容。这也是多次保存场景里最容易忽略的一处版本陷阱。1.2 它和 checkpoint、H5 到底差在哪很多人纠结保存格式我把三者放在一起对比你自己看格式存什么适合场景多次保存的风险点Checkpoint只存变量权重训练中断续训、快速落盘不带图加载要重建结构H5结构权重优化器快速原型、Keras 内部用自定义对象多时难加载SavedModel图权重签名线上部署、跨语言推理目录残留、签名覆盖结论其实很清晰训练中途要频繁保存用 Checkpoint 最省事要交付给别人做推理部署才用 SavedModel。如果你把 SavedModel 当成训练时的每步落盘手段那基本上就是在给自己挖坑因为每次保存都要重新写一遍图和变量多次保存带来的目录膨胀和残留问题会被放大。1.3 签名这个东西到底管什么签名signatures是 SavedModel 的入口约定。你保存的时候没指定签名TF 会自动帮你加一个默认的serving_default把模型的可调用部分包进去。问题是如果你反复用tf.saved_model.save覆盖同一个目录而不同次传入的signatures字典 key 不一样后一次的签名会顶掉前一次旧签名直接消失。签名里保存的是 ConcreteFunction如果这个函数闭包捕获了额外的变量或常量整个图就会被拖进saved_model.pb模型体积会明显膨胀。所以多次保存时签名不是可有可无的装饰而是每次都该显式确认的东西。2. 多次保存模型最容易踩的四个坑2.1 坑一往同一个路径反复写旧文件残留在目录里这是最高频、也最隐蔽的坑。tf.saved_model.save或model.save(path)默认行为是在目标目录里覆盖写入但它并不保证把上一版遗留的多余文件删掉。举个我自己踩过的实例第一次保存时模型变量比较多TF 自动分成了 3 个数据分片variables/ ├── variables.data-00000-of-00003 ├── variables.data-00001-of-00003 └── variables.data-00002-of-00003第二次保存时模型变小了比如剪枝之后只生成 1 个分片variables/ ├── variables.data-00000-of-00001这时候variables/目录里就同时存在着新写的00000-of-00001和旧的00001-of-00003、00002-of-00003。加载时 TF 会去读 index 文件index 指向的却是旧的00001-of-00003一读就报数据校验失败或者干脆加载出一个错误形状的权重。规避方法有两个保存前先清空目标目录最直接import shutil, os def clean_save(model, path): if os.path.exists(path): shutil.rmtree(path) model.save(path)每次保存用独立的子目录把版本信息写进路径从根上避免覆盖import os def versioned_save(model, base_dir, tag): save_dir os.path.join(base_dir, tag) os.makedirs(save_dir, exist_okTrue) model.save(save_dir) return save_dir实测下来第二种更稳因为它天然保留了历史版本回溯也方便。2.2 坑二自定义层、自定义损失加载时报 unknown class这个坑在多次保存里特别典型你在 notebook 里定义了一个MyAttentionLayer训练时保存没问题加载时却报ValueError: Unknown layer: MyAttentionLayer。原因很直白——Keras 加载 SavedModel 时需要能重新实例化你的类而这个类只存在于你当前的 Python 进程里TF 自然找不到。解决有两条路import tensorflow as tf # 方式一注册装饰器最推荐 tf.keras.utils.register_keras_serializable(packagemy_pkg) class MyAttentionLayer(tf.keras.layers.Layer): def __init__(self, units64, **kwargs): super().__init__(**kwargs) self.units units def get_config(self): config super().get_config() config.update({units: self.units}) return config# 方式二加载时显式传 custom_objects model tf.keras.models.load_model( my_model, custom_objects{MyAttentionLayer: MyAttentionLayer} )关键点在于get_config必须把自定义的构造参数都吐出来否则重建时参数丢失加载出来的模型结构和原来的不一样权重对不上会直接报 shape 不匹配。2.3 坑三签名被后一次保存覆盖线上找不到入口我在一次灰度发布里遇到过新版本模型用tf.saved_model.save保存到与旧版相同的目录结果旧版依赖的predict签名没了线上服务报Signature predict does not exist。原因是每次保存时signatures的 key 集合不同后一次会整体替换。正确做法是每次保存都显式、完整地给出签名并且保持一致import tensorflow as tf class MyModule(tf.Module): def __init__(self, model): super().__init__() self.model model tf.function(input_signature[ tf.TensorSpec(shape[None, 224, 224, 3], dtypetf.float32, nameimages) ]) def predict(self, images): return {scores: self.model(images, trainingFalse)} module MyModule(model) tf.saved_model.save(module, save_path, signatures{predict: module.predict})用tf.Module或者tf.function(input_signature...)把签名固定下来既能避免覆盖也能让导出的图更干净不至于把训练分支也塞进去。2.4 坑四训练图和推理图混着存体积翻倍还容易加载错很多人为了方便直接保存model本身结果 Keras 会把训练相关的节点dropout、BN 的 training 分支、损失函数、优化器状态全都序列化进图。多次保存后saved_model.pb的体积可能从几十 MB 涨到几百 MB而且加载时如果用错了入口推理结果会带上训练时的随机性。建议在保存推理模型时明确走一个干净的子模型或者用tf.Module包一层把trainingFalse固定进去这样每次保存产出的图都是同构的多次保存也不会互相干扰。3. 反复保存模型的正确姿势3.1 先分清tf.saved_model.save和model.save这两个 API 名字像、行为不同我列个表帮你决策API保存内容适合多次保存建议model.save(path)Keras 格式可含 H5 或 SavedModel训练完交付用带版本子目录tf.saved_model.save(obj, path)纯 SavedModel需给签名部署、跨语言显式签名清目录tf.train.Checkpoint只存变量训练中途落盘天然支持多份最安全我个人的习惯是训练循环里只写 Checkpoint训练结束时统一导一次 SavedModel。这样既避免频繁重写图也把多次保存的问题压缩到一次。3.2 目录命名策略让版本自己说话多人协作或者长期项目里我强推目录带版本戳的做法import os, time, shutil def save_with_version(model, root, prefixmodel, keep_last3): tag time.strftime(%Y%m%d-%H%M%S) save_dir os.path.join(root, f{prefix}_{tag}) os.makedirs(save_dir, exist_okTrue) model.save(save_dir) # 只保留最近 N 个防止磁盘被撑爆 versions sorted([ d for d in os.listdir(root) if d.startswith(prefix) ]) for old in versions[:-keep_last]: shutil.rmtree(os.path.join(root, old), ignore_errorsTrue) return save_dir这套逻辑解决了三个问题多版本共存、覆盖残留、磁盘膨胀。实测在 30 多天的连续训练里磁盘占用始终稳定在几个 G 以内。提示保留数量keep_last不要设太大SavedModel 单份动辄上百 MB尤其带 embedding 的模型几百 MB 很常见。3.3 加载时怎么选对签名加载 SavedModel 后第一件事是打印签名确认入口对import tensorflow as tf loaded tf.saved_model.load(model_20240101-120000) print(list(loaded.signatures.keys())) # [predict] infer loaded.signatures[predict] # 看输入输出结构避免喂错 shape print(infer.structured_input_signature) print(infer.structured_outputs)如果签名里有多个入口比如predict和embedding务必确认线上调用的是哪一个不要靠猜。多次保存场景下不同版本签名集合不一致是常态加载前先列一遍签名能省掉很多排查时间。4. 常见问题排查速查表4.1 报错信息与对应病灶报错/现象最可能的原因处理方式Data loss: not an sstable旧分片残留index 指向错文件清空目录重存Unknown layer: XXX自定义类未注册加注册装饰器或 custom_objectsSignature xxx does not exist签名被后续保存覆盖每次保存显式给同一组签名模型体积异常大训练分支被打进图用 tf.Module 包推理入口加载后推理结果随机BN/Dropout 走 training 分支固定 trainingFalse加载报 shape 不匹配结构变了但权重是老版本检查 get_config 完整返回4.2 三步定位法遇到 saved-model 加载异常我一般按下面三步走命中率很高看目录ls -R model_dir重点看variables/里分片数量是否和variables.index记录的一致多余的旧分片是重灾区。看签名加载后立刻list(loaded.signatures.keys())确认入口存在、输入 shape 和线上一致。看加载日志把 TF 日志等级调到 INFO观察它读到了哪个saved_model.pb、哪个变量文件很多问题在日志里直接写明了。注意排查时不要直接在线上目录上操作先把模型目录整个拷到本地一份避免误删线上文件。5. 一次完整的多次保存实操复盘假设我们要训练 10 个 epoch要求每个 epoch 落一次 Checkpoint第 5 个 epoch 和最后一个 epoch 各导出一次可部署的 SavedModel同时保留最近 3 个版本。完整流程如下import os import shutil import tensorflow as tf ROOT ./runs CKPT_DIR os.path.join(ROOT, ckpt) SM_DIR os.path.join(ROOT, saved_model) os.makedirs(CKPT_DIR, exist_okTrue) os.makedirs(SM_DIR, exist_okTrue) model tf.keras.Sequential([ tf.keras.layers.Input(shape(28, 28, 1)), tf.keras.layers.Conv2D(32, 3, activationrelu), tf.keras.layers.Flatten(), tf.keras.layers.Dense(10, activationsoftmax), ]) optimizer tf.keras.optimizers.Adam() ckpt tf.train.Checkpoint(modelmodel, optimizeroptimizer) ckpt_manager tf.train.CheckpointManager( ckpt, CKPT_DIR, max_to_keep3 ) def export_saved_model(model, base_dir, tag): 每次用独立子目录导出避免覆盖残留 save_dir os.path.join(base_dir, tag) if os.path.exists(save_dir): shutil.rmtree(save_dir) os.makedirs(save_dir, exist_okTrue) class ServingModule(tf.Module): def __init__(self, inner): super().__init__() self.inner inner tf.function(input_signature[ tf.TensorSpec(shape[None, 28, 28, 1], dtypetf.float32, nameimages) ]) def predict(self, images): return {probs: self.inner(images, trainingFalse)} module ServingModule(model) tf.saved_model.save( module, save_dir, signatures{predict: module.predict} ) return save_dir # 训练循环 for epoch in range(10): # 这里用假数据演示真实场景替换成你的 dataset x tf.random.normal([32, 28, 28, 1]) y tf.random.uniform([32], maxval10, dtypetf.int32) with tf.GradientTape() as tape: logits model(x, trainingTrue) loss tf.keras.losses.sparse_categorical_crossentropy(y, logits) loss tf.reduce_mean(loss) grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) ckpt_manager.save() # 每轮落一次 Checkpoint print(fepoch {epoch}, loss{loss.numpy():.4f}) if epoch 4 or epoch 9: path export_saved_model( model, SM_DIR, fv{epoch:02d} ) print(fexported - {path}) # 验证加载 loaded tf.saved_model.load(os.path.join(SM_DIR, v09)) infer loaded.signatures[predict] out infer(tf.random.normal([1, 28, 28, 1])) print(probs shape:, out[probs].shape)这段代码里有几个刻意设计的地方值得说一说Checkpoint 管训练落盘SavedModel 只管部署导出两者职责不混。多次保存的问题被彻底隔离到导出环节。export_saved_model里先rmtree再建目录杜绝了旧分片残留。这是我踩坑之后固定下来的写法。签名用tf.function(input_signature...)显式声明trainingFalse写死在函数体里保证每次导出的图同构线上加载结果稳定。用版本号做子目录多版本共存、可回溯加载时选版本目录即可。我个人在长期项目里的体会是SavedModel 本身设计没问题问题都出在人身上——图省事往一个目录反复写、不给签名、不清理历史文件。把导出环节当成一次正式的发布每次清目录、显式签名、独立版本目录这些坑基本就跟你无缘了。如果你现在正被not an sstable或者Unknown layer折磨先别急着改代码去variables/目录里数一数分片数量和文件名后缀十有八九能找到答案。
