TensorFlow模型格式转换实战:Checkpoint、SavedModel与冻结pb互转指南
1. 从一个真实场景说起为什么需要折腾模型格式转换做深度学习工程落地的朋友大概率都遇到过这种局面训练阶段用TensorFlow的Checkpoint机制保存权重训练脚本跑得好好的评估、微调都没问题可一旦要把模型交给部署团队对方开口就问“有没有pb文件”或者“能不能给个SavedModel目录”。这时候如果手头只有一堆.ckpt文件或者反过来只有SavedModel却想恢复训练就得在两种格式之间来回倒腾。我自己第一次碰到这个问题是在一个图像分类项目上。训练用的是tf.train.Saver那套老流程产出了model.ckpt-10000.index、model.ckpt-10000.data-00000-of-00001、model.ckpt-10000.meta三件套。部署同事要的是冻结后的.pb我一开始以为改个后缀就行结果当然是打不开。后来才明白Checkpoint、SavedModel、冻结图frozen graph这三者根本不是一个层面的东西它们各自承载的信息、依赖的上下文、适用的场景都不一样。这篇文章就把我这些年在这几种格式之间转换踩过的坑、总结出来的稳定流程完整梳理一遍。核心围绕三件事Checkpoint怎么转成SavedModel、SavedModel怎么转回Checkpoint、SavedModel怎么冻结成pb以及反向的pb怎么还原成可继续训练的Checkpoint。内容偏工程实操适合已经跑通过TensorFlow训练、但对模型持久化机制还比较模糊的开发者也适合需要对接部署流程的算法工程师。读完之后你应该能独立完成任意方向的格式转换并且知道每一步为什么这么做、哪里容易出错。先给一个最基础的概念对齐不然后面全是糊涂账。Checkpoint本质上是“变量字典的快照”它只存了变量名到张量值的映射外加图结构在.meta里TF2默认不存。它不包含前向计算图的完整定义也不包含签名signature所以单独一个ckpt文件是没法直接拿去推理的。SavedModel是TensorFlow官方推荐的通用序列化格式它把计算图、变量、签名、资产文件打包成一个目录语言无关、可跨平台加载是TF Serving、TF Lite转换、TF.js转换的入口。冻结图pb则是把变量全部变成常量constant后固化进GraphDef得到一个自包含的、不依赖外部变量文件的单文件模型适合嵌入式或轻量推理场景。理解了这三者的定位转换的方向感就出来了从“轻”到“重”需要补图结构和签名从“重”到“轻”需要把变量固化。下面分方向逐个拆。2. Checkpoint转SavedModel补全图结构与签名是核心难点2.1 为什么不能直接读ckpt就导出很多人第一反应是写个脚本tf.train.load_checkpoint读进来然后tf.saved_model.save出去。这条路走不通的原因在于tf.saved_model.save需要一个可追踪的Python对象通常是tf.Module或tf.keras.Model它要能对输入做一次前向追踪从而记录下完整的计算图。而Checkpoint里只有变量值没有这个可追踪对象也没有输入输出的定义。所以正确的思路是先重建模型结构再把ckpt里的权重灌进去最后导出SavedModel。重建结构这一步是绕不开的你必须知道原始模型的网络定义。这也是为什么工程上强烈建议训练时就同时保存SavedModel而不是只留ckpt。2.2 基于meta图重建TF1风格如果你的ckpt是TF1时代用tf.train.Saver保存的通常会带一个.meta文件里面存了GraphDef。这种情况下可以用tf.train.import_meta_graph把图恢复出来再restore变量import tensorflow as tf meta_path ./model.ckpt-10000.meta ckpt_prefix ./model.ckpt-10000 saver tf.train.import_meta_graph(meta_path, clear_devicesTrue) graph tf.get_default_graph() with tf.Session(graphgraph) as sess: saver.restore(sess, ckpt_prefix) # 此时graph里已经有完整的计算图可以找到输入输出tensor input_tensor graph.get_tensor_by_name(input:0) output_tensor graph.get_tensor_by_name(logits:0)拿到图之后导出SavedModel需要构造签名。TF1里常用tf.saved_model.simple_savetf.saved_model.simple_save( sess, ./saved_model_dir, inputs{input: input_tensor}, outputs{output: output_tensor} )这里有个关键点get_tensor_by_name里的名字必须和建图时一致。如果你不记得tensor名字可以在restore之后遍历graph.get_operations()打印出来找。我一般会写个小工具函数把所有placeholder和输出节点列出来避免瞎猜。2.3 基于代码重建TF2/Keras风格TF2默认用tf.kerasckpt往往是通过model.save_weights或tf.train.Checkpoint保存的没有meta图。这时候唯一可靠的办法是用原始模型定义代码重新实例化一个模型然后load_weights或restoreimport tensorflow as tf # 必须和训练时完全一致的模型定义 model build_model() # 你自己定义的函数 model.load_weights(./checkpoints/ckpt-10000) # 直接导出SavedModel model.save(./saved_model_dir, save_formattf)model.save在TF2里默认就是SavedModel格式会自动带上签名。如果模型有自定义的call逻辑建议显式指定signatures参数否则默认签名可能只接受特定shape的输入。注意load_weights对变量名的匹配非常严格。如果训练时用了tf.train.Checkpoint包装恢复时也要用同样的包装结构否则会出现大量“未匹配变量”的警告甚至静默加载失败。我踩过一次坑模型结构看着一样但因为一层Dense的命名从dense变成了dense_1权重全没加载上推理结果完全是随机的。2.4 变量名不匹配的排查方法变量名对不上是ckpt转SavedModel最常见的拦路虎。排查手段有两个一是用tf.train.list_variables(ckpt_path)列出ckpt里所有变量名和shape二是用model.variables或model.weights列出当前模型的变量名。两边打印出来对比差异一目了然。import tensorflow as tf ckpt_vars tf.train.list_variables(./checkpoints/ckpt-10000) for name, shape in ckpt_vars: print(name, shape) for v in model.variables: print(v.name, v.shape)如果只是命名前缀不同比如多了model/或layer_with_weights-0/可以用tf.train.Checkpoint的save/restore配合assert_consumed来定位或者手动做名字映射。手动映射虽然笨但在结构复杂时最可靠。3. SavedModel转Checkpoint恢复训练与继续微调3.1 什么情况下需要反向转换SavedModel转Checkpoint的需求通常出现在两个场景一是拿到别人训练好的SavedModel想在自己的数据上继续微调二是部署环境只给了SavedModel但研发需要恢复训练流程做实验。这时候需要把SavedModel里的变量提取出来重新组织成Checkpoint格式。3.2 加载SavedModel并提取变量TF2里加载SavedModel很简单import tensorflow as tf loaded tf.saved_model.load(./saved_model_dir) # loaded是一个可调用对象内部变量可以通过variables访问 print(list(loaded.variables))但要注意loaded.variables返回的是变量列表顺序和名字不一定和原始训练时一致。更稳妥的方式是通过loaded.signatures找到具体的计算函数再顺着函数追踪变量。不过对于大多数Keras导出的SavedModel直接用tf.keras.models.load_model加载会更友好model tf.keras.models.load_model(./saved_model_dir) model.save_weights(./new_ckpt/ckpt)这样得到的Checkpoint就是标准的TF2权重格式可以直接被后续训练脚本用load_weights读取。3.3 用tf.train.Checkpoint重新包装如果后续训练脚本用的是tf.train.Checkpoint管理状态比如包含优化器状态、步数等那就需要手动构造一个Checkpoint对象import tensorflow as tf model tf.keras.models.load_model(./saved_model_dir) optimizer tf.keras.optimizers.Adam() ckpt tf.train.Checkpoint(modelmodel, optimizeroptimizer) ckpt.write(./new_ckpt/ckpt)这样写出来的Checkpoint包含了模型和优化器的变量。恢复时用同样的结构ckpt.restore即可。这里有个细节优化器状态在SavedModel里通常是不保存的所以反向转换出来的Checkpoint一般只含模型权重优化器状态需要重新初始化。如果要做无缝续训最好在训练阶段就定期保存完整的tf.train.Checkpoint。3.4 转换后验证权重是否真的对上了转换完不做验证等于白转。我的习惯是转换后立刻用同一份输入分别跑原SavedModel和新Checkpoint恢复的模型对比输出是否一致。import numpy as np sample np.random.rand(1, 224, 224, 3).astype(np.float32) # 原SavedModel输出 orig_out model(sample, trainingFalse).numpy() # 新Checkpoint恢复后输出 new_model build_model() new_model.load_weights(./new_ckpt/ckpt) new_out new_model(sample, trainingFalse).numpy() print(np.max(np.abs(orig_out - new_out)))理想情况下这个差值应该是0或者极小浮点误差级别。如果差值很大说明权重没对上得回去查变量名映射。4. SavedModel冻结成pb把变量钉死成常量4.1 冻结的本质与适用边界冻结freeze做的事情是把计算图里所有Variable节点替换成Const节点值就是当前变量的取值。这样得到的GraphDef不再依赖任何外部变量文件单文件即可推理。它适合的场景是模型结构固定、不需要再训练、部署环境希望尽量简单。但冻结也有代价模型体积会变大因为常量直接嵌在图里没有压缩而且失去了灵活性改不了输入shape、加不了新层。所以如果你的部署环境支持SavedModel优先用SavedModel冻结pb更多是历史遗留或特定嵌入式场景的需求。4.2 TF1风格的freeze_graph流程TF1时代标准做法是用tf.compat.v1.graph_util.convert_variables_to_constantsimport tensorflow as tf with tf.Session() as sess: saver tf.train.import_meta_graph(./model.ckpt-10000.meta) saver.restore(sess, ./model.ckpt-10000) graph tf.get_default_graph() input_tensor graph.get_tensor_by_name(input:0) output_tensor graph.get_tensor_by_name(logits:0) frozen_graph tf.compat.v1.graph_util.convert_variables_to_constants( sess, graph.as_graph_def(), [output_tensor.op.name] ) with tf.io.gfile.GFile(./frozen_model.pb, wb) as f: f.write(frozen_graph.SerializeToString())convert_variables_to_constants的第三个参数是输出节点名字列表必须准确。如果写错冻结出来的图可能缺少关键节点加载时报“找不到节点”。4.3 TF2里从SavedModel冻结TF2没有直接的冻结API但可以用convert_variables_to_constants_v2import tensorflow as tf from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 model tf.keras.models.load_model(./saved_model_dir) # 构造一个具体的函数 full_model tf.function(lambda x: model(x)) full_model full_model.get_concrete_function( tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype) ) frozen_func convert_variables_to_constants_v2(full_model) frozen_func.graph.as_graph_def() tf.io.write_graph( frozen_func.graph, ./, frozen_model.pb, as_textFalse )这段代码的关键是get_concrete_function它把动态的Keras模型固定成一个具体输入shape的计算图。输入shape必须和部署时一致否则冻结出来的图只能接受那个特定shape。4.4 冻结后如何验证pb可用冻结完别急着交付先用tf.compat.v1.GraphDef读回来确认节点齐全import tensorflow as tf with tf.io.gfile.GFile(./frozen_model.pb, rb) as f: graph_def tf.compat.v1.GraphDef() graph_def.ParseFromString(f.read()) with tf.Graph().as_default() as graph: tf.import_graph_def(graph_def, name) for op in graph.get_operations(): print(op.name, op.type)重点看输入输出节点在不在以及有没有残留的VariableV2或Variable类型节点。如果还有变量节点说明冻结没成功通常是输出节点列表没覆盖到所有分支。5. pb还原成Checkpoint逆向操作与它的局限5.1 pb能不能还原成可训练模型严格来说冻结pb无法完美还原成可训练的Checkpoint因为变量已经变成常量常量是不可训练的。但如果你只是想把pb里的权重提取出来重新灌到一个结构相同的模型里这条路是通的。做法是解析pb的GraphDef找到所有Const节点按名字映射回模型的变量。5.2 从pb提取常量并映射回变量import tensorflow as tf import numpy as np with tf.io.gfile.GFile(./frozen_model.pb, rb) as f: graph_def tf.compat.v1.GraphDef() graph_def.ParseFromString(f.read()) const_values {} for node in graph_def.node: if node.op Const: tensor node.attr[value].tensor arr tf.make_ndarray(tensor) const_values[node.name] arr # 然后按名字匹配到你的模型变量 model build_model() for var in model.variables: key var.name.split(:)[0] if key in const_values: var.assign(const_values[key])这段代码的难点在于名字匹配。冻结时节点名可能被加了前缀或做了重命名需要根据实际情况调整。我一般会先把const_values的key全部打印出来和model.variables的名字对照找出规律再写映射。5.3 还原后的训练可行性评估即使权重成功灌回去了也要清楚这个模型是从冻结图来的可能丢失了训练时的很多信息比如BN层的移动均值/方差是否完整、dropout配置是否正确、损失函数定义是否还在。所以还原后第一件事是跑一遍评估确认精度和原模型一致再决定要不要继续训练。如果精度对不上说明某些关键变量没还原到位。6. 踩坑实录那些让我加班到深夜的转换问题6.1 输入shape写死导致的部署失败有一次我把SavedModel冻结成pb输入shape写的是[1, 224, 224, 3]结果部署方要批量推理喂了[8, 224, 224, 3]直接报错。冻结时如果用get_concrete_function指定了固定batch图里就钉死了。解决办法是冻结时把batch维度设为Nonefull_model full_model.get_concrete_function( tf.TensorSpec([None, 224, 224, 3], tf.float32) )这样冻结出来的pb支持动态batch。但要注意有些算子对动态shape支持不好冻结后可能仍然报错需要实测。6.2 BN层与dropout在冻结时的行为差异BatchNormalization和Dropout在训练和推理时行为不同。冻结时必须确保模型处于推理模式trainingFalse否则BN会用当前batch的统计量dropout会随机丢弃冻结出来的结果就是错的。Keras里model.save和convert_variables_to_constants_v2默认会处理这个问题但如果你手动构造tf.function一定要显式传trainingFalse。6.3 变量名带作用域前缀的匹配问题TF1的变量名经常带scope/前缀比如conv1/weights。而Keras模型的变量名可能是conv2d/kernel。这种命名体系差异导致直接按名字匹配几乎不可能。我的经验是如果两边都是自己可控的代码训练时就统一用Keras风格如果是接手别人的模型就老老实实写映射表一个变量一个变量对。6.4 转换后精度掉点的排查思路转换后精度掉点排查顺序建议是先确认输入预处理是否一致归一化参数、通道顺序再确认模型是否处于推理模式然后对比中间层输出。我遇到过一次是pb里的输入节点名字变了部署方喂错了tensor数据根本没进网络。所以转换后一定要做端到端的数值对比不能只看文件生成成功就完事。7. 一套可复用的转换检查清单把上面这些经验固化成一个检查清单每次转换照着过一遍能省掉大量返工。检查项具体内容常见问题模型结构一致性转换前后模型定义是否完全一致层命名不同导致权重加载失败变量名匹配ckpt变量名与模型变量名逐一对照作用域前缀差异输入输出节点确认tensor名字和shape名字写错、shape写死推理模式BN/Dropout是否处于推理状态忘记传trainingFalse数值验证同一输入对比转换前后输出精度掉点、输出全随机动态shape是否需要支持变长/变batch输入冻结时写死shape优化器状态续训场景是否需要优化器变量SavedModel不含优化器状态这份清单看着简单但每一条背后都是实打实的坑。尤其是变量名匹配和数值验证这两项我建议做成自动化脚本每次转换后自动跑一遍比人眼靠谱得多。8. 关于工具链选择的一点个人看法现在TensorFlow生态里做格式转换工具其实不少但各有各的脾气。tf.saved_model.save和tf.keras.models.load_model是最省心的前提是模型用Keras标准写法。一旦涉及自定义层、自定义训练循环就得自己处理签名和变量追踪复杂度陡增。convert_variables_to_constants_v2虽然能用但它是内部API版本之间行为可能有差异升级TF版本时要重点回归测试。我的建议是新项目一律以SavedModel为唯一交付格式训练时定期保存tf.train.Checkpoint用于续训部署时从SavedModel导出所需格式。Checkpoint和pb之间的直接转换尽量少做因为中间丢失的信息太多每次转换都是一次信息损耗。如果非要转就把转换脚本和验证脚本一起维护当成项目的一部分而不是临时写个一次性脚本用完就扔。最后分享一个我用了很久的小技巧在训练脚本里加一个export_saved_model的回调每隔几个epoch就导出一份SavedModel。这样即使训练中断手头也永远有一份可直接部署的模型不用事后从ckpt里费劲重建。这个习惯帮我省过好几次急活推荐你也试试。