zoukankan      html  css  js  c++  java
  • 合并多个tensorflow模型的办法

    直接上代码:

    import tensorflow as tf
    from tensorflow.python.tools import freeze_graph
    from tensorflow.python.framework.graph_util import convert_variables_to_constants
    import os
    import numpy as np
    
    filename1 = "model_a.pb"
    filename2 = "model_b.pb"
    
    def load_graphdef(filename):
    	with tf.gfile.GFile(filename, "rb") as f:
    		graph_def = tf.GraphDef()
    		graph_def.ParseFromString(f.read())
    	return graph_def
    
    
    
    def load_graph(graph_def, prefix):
    
    	with tf.Graph().as_default() as graph:
    		tf.import_graph_def(graph_def, name=prefix) 
    	
    	return graph
    
    graph1 = load_graphdef(filename1)
    graph2 = load_graphdef(filename2)
    
    graph1_out, = tf.import_graph_def(graph1, return_elements=['mode_a_output:0'], name="model_a")
    graph2_out, = tf.import_graph_def(graph2, return_elements=['mode_b_output:0'], name="model_b")
    
    
    z = tf.concat([graph1_out,  graph2_out], 1)
    
    tf.identity(z, "merge_output")	
    
    init_op = tf.global_variables_initializer()
    with tf.Session() as sess:
    	sess.run(init_op)
    	graph = convert_variables_to_constants(sess, sess.graph_def, ["merge_output"])
    	tf.train.write_graph(graph, '.', 'merge.pb', as_text=False)
    

    合并后的pb文件,输入节点为原来输入节点的并集。和原模型输入的区别是:输入节点分别增加的对应的前缀model_a/, model_b/。

  • 相关阅读:
    软件工程学习总结
    第13次作业--邮箱的正则表达式
    第12次作业--你的生日
    第11次作业--字符串处理
    第10次作业
    找回感觉的练习
    基础网络流学习笔记
    卷积定理的证明
    快速xxx变换相关
    主席树相关
  • 原文地址:https://www.cnblogs.com/th3Bear/p/11438310.html
Copyright © 2011-2022 走看看