首页 > 其他 > 详细

合并多个tensorflow模型的办法

时间:2019-08-31 11:52:16      阅读:156      评论:0      收藏:0      [点我收藏+]

直接上代码:

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/。

合并多个tensorflow模型的办法

原文:https://www.cnblogs.com/th3Bear/p/11438310.html

(0)
(0)
   
举报
评论 一句话评论(0
关于我们 - 联系我们 - 留言反馈 - 联系我们:wmxa8@hotmail.com
© 2014 bubuko.com 版权所有
打开技术之扣,分享程序人生!