Question

我有一个* .pb格式的预训练张量流模型。我也有一个带有特定输入和输出网络的管道。现在，我想使用预训练的模型将其加载到我的管道中：

# inside pipeline.py

tf.reset_default_graph()
graph1 = tf.Graph()

# Building my input:

input_tensor = tf.placeholder('float', shape=input_shape, name='image')
# ....

# Now I want to use features from a pretrained model:
from ABCModel import model
net = model.build('modelABC.pb', input_tensor)

# Next, I use some features from this model to construct a new output, e.g.:

output = tf.contrib.layers.fully_connected(['feature_tensor_A'], 100)

在ABCModel.py内部，我有一个构建函数：

def build(pb_file, input_tensor):

    f = gfile.FastGFile(pb_file, 'rb')
    graph_def = tf.GraphDef()
    # Parses a serialized binary message into the current message.
    graph_def.ParseFromString(f.read())
    f.close()

    _ = tf.import_graph_def(graph_def , input_map={'input0': input_tensor})

问题是，如果仅执行此操作，导入的模型张量就不会成为graph1的一部分，我设法通过在pipeline.py的开头添加以下内容来使它们可用：

cm = graph1.as_default()
cm.__enter__()

但是，这似乎大大降低了梯度计算的速度。在这样的管道内导入预训练图的更有效，更优雅的方法是什么？

后续问题：我已将我的graph1设置为默认图，但是当我从graph_def导入时，创建了一个新图，或者这是什么问题？在文档中，它说该图已添加到现有的默认图（？）中-但这似乎不起作用。

使用graph_def和输入映射构建新的Tensorflow模型

0 个答案: