我做统一推理服务的时候遇到一个很典型的需求一个进程里同时加载多个TensorFlow模型。之前单模型推理一直没问题就是tf.Session()一开、模型一load、完事。直到产品那边说要同时上一个图像检测模型和一个文本分类模型我天真的以为把加载代码复制两遍就行结果连环踩坑先是Graph节点冲突再是显存直接OOM最后发现Session和Graph的绑定关系也没搞对。折腾了一个多星期才把一套能稳定撑住多模型并存的方案跑通。这个系列已经写到第3篇了前两篇我都在讲单模型的训练和导出这回来聊聊多模型同时加载。这个需求在真实项目里太常见了统一推理服务要承载多个模型、A/B实验需要把两个模型同时拉起来对比、或者一个业务逻辑里要先过检测模型再过分类模型。多模型加载的难点不在于API语法而在于很多人和我一样一开始根本没想到TensorFlow默认会把所有东西都塞进同一张全局计算图里模型一多就开始互相抢地盘。这篇我用自己从翻车到稳定运行的完整过程把多模型加载的核心原理、隔离方案、三种模型格式的加载差异、显存控制和并发调优讲清楚。如果你正准备在同一个进程里跑多个TensorFlow模型可以直接照着文末的模型管理器代码来能少走很多弯路。1. 同时加载多个模型最容易踩的暗坑在哪1.1 默认只有一张全局计算图模型一多就互相抢地盘TensorFlow里的Graph计算图是所有变量、占位符、常量和算子的容器。如果不主动新建Graph所有操作都会被注册到一张全局默认图上。单模型场景下这个设计没有存在感但是两个模型一起加载的时候问题一下就暴露了。我第一次把模型A和模型B的加载代码前后排列第二个模型加载直接报错。原因很典型两个模型内部都有input、output、conv1这种通用节点名TensorFlow在第二张图插入同名节点时自动把它重命名为input_1、conv1_1。表面上看程序没崩但后续如果你用原始节点名去拿张量拿到的可能是完全错误的内容。如果两个模型结构完全一样重名变量甚至会直接抛Variable already exists异常。还有一个更隐蔽的坑Graph在Session创建之后会被自动冻结finalized也就是不能再往里面添加任何节点了。我先加载模型A、创建Session然后再加载模型B这时候模型B的节点想加进同一个Graph直接报Cannot add more ops to graph。这个机制本身是为了防止训练时Session和Graph状态不一致但在多模型加载场景里它变成了一个必须先绕过去的坎。1.2 Session如果不指定Graph会悄悄绑定当前默认图Session是TensorFlow执行模型的环境承载着运行时状态和计算资源。创建一个Session时如果不传graph参数它就会绑定当前默认图。很多人写多模型加载时其实已经创建了多个Session对象以为各自独立了但节点全部都在同一张默认图里你根本没法区分哪个节点属于哪个模型。更麻烦的是Session一旦创建它绑定的Graph就被冻住了后续另一个模型想往这张图里加节点是加不进去的。反过来操作——先把两个模型都加载完、再创建Session——情况会更乱两个模型的节点混在同一张图里变量名全靠系统自动改名撑着模型A的Session甚至能run到模型B的节点这种代码别说维护了线上排查都会让你想砸键盘。我自己在实际定位时吃过一次大亏模型A和模型B明明都加载成功了但推理时输入张量关键字的形状对不上查了半天才发现是自动改名把张量索引带偏了。从那时起我就确认了一个结论——多模型加载必须做图级别隔离。1.3 显存和线程调度上的隐形冲突多模型同时加载不光有图层面的问题资源层面的冲突同样致命。TensorFlow在加载一个Session时默认会预占大量显存如果第一个Session创建时把GPU上可用显存几乎全占完了第二个Session加载必然OOM。我一开始就是先加模型A显存占用飙升再加模型B时直接报Resource exhausted。CPU线程资源也一样。每个Session会按照配置创建自己的线程池多个Session各自单独看都没问题放到同一台机器上就会互相争抢CPU核。模型少的时候还能忍模型一多就会出现请求抖动、推理耗时忽高忽低的情况。这些问题不属于加载API层面但只要你打算做多模型混部迟早会撞上。2. 核心方案每个模型一套SessionGraph彻底隔离2.1 为什么我放弃图融合而选隔离有些资料会建议把多个模型融合成一张大图把模型A的权重固化后通过改图把输出直接接到模型B的输入上这样理论上能省掉重复的框架开销。这个思路在端侧部署或极简网络等特定场景下是有价值的但在服务端多模型混部场景我并不推荐优先走这条路。第一个原因是模型内部结构往往由训练脚本或导出的工具生成手工改图非常脆弱第二个原因是不同模型迭代节奏不一样今天改模型A、明天改模型B融合图的每次升级都要整体重建第三个原因是隔离方案的语义足够清晰哪个模型出了问题直接定位到对应的Session和Graph不需要在一堆自动改名的节点里大海捞针。当然同进程隔离也不是唯一方案。多进程隔离更彻底一个模型一个进程崩了也不影响其他模型但进程资源开销大、进程间数据传输麻烦还得考虑共享内存或RPC对大多数业务场景来说太重了。综合权衡下来我的选择是同进程内做每个模型一套Session 一套Graph的硬隔离用管理器统一维护。2.2 模型管理器把加载和预测封装成一张表具体实现上我把模型管理封装成一个Python类核心思路是每个模型创建独立的Graph创建Session时显式传入这个Graph后续所有加载和推理都在各自Graph上下文中进行。用名字作为索引维护一份模型注册表。import tensorflow as tf import threading class MultiModelServer(object): def __init__(self): self.sessions {} self.graphs {} self.inputs {} self.outputs {} self.locks {} def add_saved_model(self, name, export_dir, tagsNone, signature_keyserving_default): tags tags or [serve] graph tf.Graph() with graph.as_default(): sess tf.Session(graphgraph) meta_graph_def tf.saved_model.loader.load( sess, tags, export_dir) signature meta_graph_def.signature_def[signature_key] inputs { key: graph.get_tensor_by_name(ts.name) for key, ts in signature.inputs.items() } outputs [ graph.get_tensor_by_name(ts.name) for ts in signature.outputs.values() ] self.sessions[name] sess self.graphs[name] graph self.inputs[name] inputs self.outputs[name] outputs self.locks[name] threading.Lock() def predict(self, name, data_dict): inputs self.inputs[name] feed_dict {inputs[key]: data_dict[key] for key in data_dict} with self.locks[name]: return self.sessions[name].run( self.outputs[name], feed_dictfeed_dict)使用的时候每个模型调用一次add_saved_model预测时把模型名和数据传进去就行。这里最关键的是tf.Session(graphgraph)这一句必须显式指定graph。如果不指定Session默认绑定的是全局默认图之前说的所有问题都会原样复发。2.3 为什么Session创建也要放在Graph上下文里有个细节我第一次写的时候没注意建Graph对象是一回事模型加载进来的节点落在哪张图是另一回事。tf.saved_model.loader.load执行时会把它解析出来的节点追加到当前默认图上。如果你在使用with graph.as_default()这个上下文管理器的代码块外面调用load那么节点会全部注册到全局默认图你的自定义Graph就成了一张空图后面get_tensor_by_name会直接报KeyError: The name ... refers to a Tensor which does not exist。所以代码里我把Session创建和模型加载都放进了with graph.as_default():块内。这样能保证三样东西始终绑定在一起模型节点在哪个Graph上Session绑定的是哪个Graph后续取张量从哪个Graph里取。只要这个铁三角关系不裂开多个模型之间就不会互相干扰。2.4 别忘了给每个模型加一把锁模型管理器的代码里我给每个模型配了一个threading.Lock。这不是多此一举。多个请求同时打同一个Session时虽然我在实际压测中很少遇到直接崩溃的情况但偶尔会出现非常诡异的卡顿加上多个Session同时执行时GPU显存申请和释放的竞争问题会变得更加隐蔽。TensorFlow官方对Session并发执行的说明比较保守我也见过有人在多线程环境里让同一个Session反复执行导致的张量状态串扰案例。稳妥的做法就是加一个模型级别的互斥锁让每次推理串行化。单模型并发能力不够的时候不要急着去掉锁更推荐的做法是给这个模型多建几个Session副本每个副本服务一组请求从横向扩容的角度去解决问题。毕竟同一个模型多Session副本共享的是同一份参数文件分别加载到各自的Graph和Session里内存占用会高一点但并发安全系数要高得多。3. 三种常见模型格式的加载差异SavedModel、Checkpoint、Frozen PB3.1 SavedModel加载生产环境最推荐的方式上面代码里用的是tf.saved_model.loader.load加载的是TensorFlow SavedModel格式。这个格式包含完整的模型结构、变量、签名和资产文件是一个自包含的目录结构部署时不需要再去重建模型网络结构对生产环境特别友好。SavedModel里有个核心概念叫SignatureDef。它定义了一个服务的输入输出协议。我上面的代码从signature_def[signature_key]里读取inputs和outputs然后通过graph.get_tensor_by_name拿到对应的Tensor对象。签名只要设计好模型内部怎么改对外服务接口都不变这对多模型统一管理帮助很大。多模型加载时我强烈建议生产环境统一用SavedModel。它同时解决了结构恢复和输入输出协议的问题Graph隔离后各模型之间互相完全透明。模型版本迭代时只需要把新的SavedModel目录替换进去管理器的add_saved_model逻辑一行都不用改。3.2 Checkpoint加载必须先重建图再恢复权重如果你手上的模型是checkpoint文件加载方式就完全不同了。Checkpoint只保存了变量值不带网络结构所以恢复前必须先把网络结构重建一遍。graph tf.Graph() with graph.as_default(): model build_model() # 重新搭建网络结构 saver tf.train.Saver() sess tf.Session(graphgraph) saver.restore(sess, checkpoint_path)这里要特别注意一点build_model()创建的所有变量必须是在当前这个graph.as_default()上下文中完成的否则Saver恢复的时候会找不到对应变量。如果两个模型结构相同变量名冲突的概率极高最稳妥的方案是在构建模型网络时用tf.variable_scope给每个模型加一个独立前缀比如with tf.variable_scope(model_a): ...这样两个模型即使结构完全一样变量也能各自存在于不同的命名空间里。3.3 Frozen PB加载用import_graph_def把图导入到指定GraphFrozen Graph也叫frozen pb是把模型结构和变量权重固化成一个GraphDef文件部署最轻量。加载它常用tf.import_graph_def这个函数可以把整个GraphDef导入到当前默认图里所以隔离方式也简单graph tf.Graph() with graph.as_default(): graph_def tf.GraphDef() with open(frozen_model_path, rb) as f: graph_def.ParseFromString(f.read()) tf.import_graph_def(graph_def, namemodel_b) sess tf.Session(graphgraph) input_tensor graph.get_tensor_by_name(model_b/input:0) output_tensor graph.get_tensor_by_name(model_b/output:0)注意看tf.import_graph_def里的name参数我给模型B的所有节点加了model_b前缀。这个前缀在加载时就规避了节点重名问题后面取张量也需要把前缀带上。如果你的pb文件里节点本来就是通用名字比如input和output多个模型同时加载时不加前缀绝对会翻车。3.4 三种格式怎么选我的经验是别在部署阶段做重构如果让我给一个排序生产环境首选SavedModel其次Frozen PBCheckpoint只建议在训练脚本内部做快速验证时用。SavedModel自带签名和版本信息多模型管理的代码也最简单Frozen PB需要手动管理节点前缀Checkpoint要重建结构在多模型场景下是最容易出问题的。你手上如果只有checkpoint文件也别急着在部署代码里现搭网络结构做恢复我建议先离线把checkpoint导成SavedModel部署那边永远只认SavedModel这一种格式。这样模型训练、导出、上线三条链路各自独立多模型加载的复杂度会被压到最低。4. 显存控制与并发推理上线前必须调好4.1 GPU显存的两个关键配置多模型同时加载显存配置不提前做好第二个Session基本就是OOM的下场。TensorFlow里常用的两个配置是allow_growth和per_process_gpu_memory_fraction。配置项行为适用场景gpu_options.allow_growthTrue显存按需增长用多少申请多少多模型混部显存分配灵活per_process_gpu_memory_fraction0.3固定预占该比例显存需要精确控制单个模型显存上限我的实际用法是如果GPU上同时跑两个模型给每个Session单独配置per_process_gpu_memory_fraction比如模型A固定40%、模型B固定30%加起来留出余量防止任何一个模型把显存吃穿。如果是简单场景开allow_growthTrue也能跑但要注意显存增长到峰值之后不会自动降回去长时间运行还是会有碎片化问题。config tf.ConfigProto() config.gpu_options.allow_growth True sess tf.Session(graphgraph, configconfig)多模型加载还有一个容易被忽略的部署技巧通过环境变量CUDA_VISIBLE_DEVICES可以让不同模型天然落在不同GPU上。模型A只看到GPU0模型B只看到GPU1显存完全隔离连per_process_gpu_memory_fraction都不用调。缺点是GPU数量有限模型多了就不够分。4.2 并发推理时的线程池配置Session默认会按CPU核数创建线程池多个Session叠加起来线程数可能远超CPU核心数。我遇到过并发一上来CPU利用率乱跳、推理耗时抖动的情况最后就是通过限制每个Session的线程数解决的。config tf.ConfigProto( intra_op_parallelism_threads4, inter_op_parallelism_threads4, )intra_op控制单个算子内部的并行线程数inter_op控制不同算子之间的并行线程数。对于推理任务这两个值不需要设很大适合多模型混部场景。具体数值还是得结合模型的计算密度和现场压测来定我给一个初步参考值单模型推理服务设成4~8之间多模型混部时整体线程数不要超过物理核数太多。4.3 动态加载、卸载和热更新多模型服务还有一个常见需求是动态加模型或下线模型。我的管理器直接在注册表里操作add_saved_model是动态加卸载时把Session close掉从各个dict里删掉对应键就行。不过有个坑要提醒Session close之后显存不一定立刻全部归还给操作系统尤其是TensorFlow显存分配器会在进程内缓存一部分。我在长期运行的服务里遇到过显存只增不减的情况排查下来是旧Session对象的引用没释放干净。卸载模型时除了close要把注册表里的所有引用都删掉必要时可以调一下gc.collect()。如果模型迭代非常频繁我倾向于直接走容器级重启让显存彻底归零虽然粗暴但最可靠。5. 排错速查与几点维护心得5.1 多模型加载常见错误快查表报错信息大概率原因处理思路Tensor ... is not an element of this graph取张量用的Graph和模型加载用的Graph不是同一个统一在graph.as_default()上下文中取张量Cannot add more ops to graphSession已经创建Graph被冻结还在往同一张图加节点每个模型单独建Graph再建SessionVariable ... already exists多个模型在同一张默认图上变量重名用variable_scope给模型加前缀或者做Graph隔离Resource exhausted: OOM / failed to allocate memory显存被第一份Session全部预占配置allow_growth或per_process_gpu_memory_fraction推理请求偶尔卡死或超时同一个Session并发run不稳定加模型级锁或拆多个Session副本这几条我在调试过程中反复遇到过基本覆盖了多模型加载初期遇到的大部分问题。5.2 调试技巧和日志避坑多模型调试最大的痛点是不知道某个节点到底在哪张图上。我建议在加载完每个模型后打印一下这个Graph里的节点名列表用命名前缀做区分一旦发现两个模型的节点混在一起代码显然在物理上就没隔离干净。node_names [node.name for node in graph.as_graph_def().node] print(node_names[:20])日志方面TensorFlow默认会刷一大堆训练和加载日志多模型时两倍刷屏。可以设置环境变量TF_CPP_MIN_LOG_LEVEL2过滤掉INFO级别的输出只保留WARNING和ERROR。排查真正的问题时再临时调回默认级别能帮你快速过滤噪音。Graph可视化也是一个好手段把每个Graph的GraphDef导出来在TensorBoard里看结构命名前缀一眼就能看出模型边界是否清晰。我每次上线新模型都会看一眼图结构确认只包含该模型自己的节点才放心压测。5.3 长期维护下来的一点经验说了这么多技术细节最后分享一个我在实际维护中体会最深的经验多模型加载跑通之后最难的不是技术实现而是后续的命名管理。每个模型从一开始就要有一个全局唯一的短名字加载时用这个名字作为Graph隔离的key日志输出带这个名字监控指标也带这个名字TensorBoard图里用这个名字做前缀。这样模型一多你看到日志就能知道是哪个模型出问题看到图就知道哪些节点属于谁。还有一点不要盲目相信我代码里写了独立Graph就一定隔离了。我多次遇到实际节点落在默认图里的情况每次排查的最终手段都是打印节点名列表。代码写得再漂亮都不如现场输出的节点列表可靠。多模型加载这件事本质上就是把每个模型一个独立计算环境这个约束贯彻到从加载、取张量、推理、日志到监控的全链路里哪个环节漏掉了它就会在哪个环节出问题。