向 Tensorflow freeze_graph 添加新層？

Question

這些討論討論了 ( 1 , 2 ) 關於向 Tensorflow 圖添加新層並重新訓練模型的問題。

下面的代碼顯示了向恢復的可訓練模型添加新層。

import tensorflow as tf

sess=tf.Session()    
#First let's load meta graph and restore weights
saver = tf.train.import_meta_graph('my_test_model-1000.meta')
saver.restore(sess,tf.train.latest_checkpoint('./'))


# Now, let's access and create placeholders variables and
# create feed-dict to feed new data

graph = tf.get_default_graph()
w1 = graph.get_tensor_by_name("w1:0")
w2 = graph.get_tensor_by_name("w2:0")
feed_dict ={w1:13.0,w2:17.0}

#Now, access the op that you want to run. 
op_to_restore = graph.get_tensor_by_name("op_to_restore:0")

#Add more to the current graph
add_on_op = tf.multiply(op_to_restore,2)

print sess.run(add_on_op,feed_dict)
#This will print 120.

但我喜歡添加圖層來恢復凍結圖。

我只為一個應用程序凍結了模型。 我喜歡在模型中添加圖層並再次凍結。 這些層更多用於后期處理，不需要訓練，因此不在訓練模型中。

我將凍結圖轉換為 TensorRT 的原因是我喜歡將這些層包含到 Int8 引擎中。

Answer 1

希望下面能幫到你。 我有一個自定義 Op，它應該添加到我從 .pb 文件（凍結模型文件）加載的現有圖形中，這樣我就能夠將新節點附加到我現有的圖形中。

Source code below: 

import tensorflow as tf
from tensorflow.python.framework import load_library
from tensorflow.python.platform import resource_loader

from tensorflow.core.protobuf import saved_model_pb2
from tensorflow.python.util import compat


# Utility functions for Loading and Freezing graphs


def load_graph(frozen_graph_filename):

    with tf.gfile.GFile(frozen_graph_filename, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name="")

    return graph

def freeze_graph(sess, output_graph):

    output_node_names = [
        "custom_op_zero","custom_op_zero_1"
  output_node_names = ",".join(output_node_names)

    output_graph_def = tf.graph_util.convert_variables_to_constants(
        sess,
        tf.get_default_graph().as_graph_def(),
        output_node_names.split(",")
    )

    with tf.gfile.GFile(output_graph, "wb") as f:
        f.write(output_graph_def.SerializeToString())

    print("{} ops written to {}.".format(len(output_graph_def.node), output_graph))


## load custom Ops shared object file

zero_out_ops = load_library.load_op_library(
    resource_loader.get_path_to_datafile('my-op/tensorflow_zero_out/python/ops/_zero_out_ops.so'))
zero_out = zero_out_ops.zero_out

frozen_graph = load_graph("frozen_model.pb")
all_tensors = [tensor for op in frozen_graph.get_operations() for tensor in op.values()]
#print (all_tensors[29])

# Input to the new node is the output of last node

zero_out_custom = zero_out(all_tensors[-1],name="custom_op_zero")
zero_out_custom1 = zero_out(all_tensors[-1],name="custom_op_zero_1")
#print (new_op)

# save new freezed model file
with tf.Session(graph=frozen_graph) as persisted_sess:
  for op in persisted_sess.graph.get_operations():
     print(op)
  freeze_graph(persisted_sess,"new_model.pb")