时间戳

首页 > 解决方案 > 如何使用 tensorflow .pb 文件？

tensorflow - 如何使用 tensorflow .pb 文件？

问题描述

我有一个 Tensorflow 文件AlexNet.pb，我正在尝试加载它，然后对我拥有的图像进行分类。我已经搜索了几个小时，但我仍然找不到加载它然后对图像进行分类的方法。是否如此明显，我只是如此愚蠢，因为似乎没有人有加载和运行 .pb 文件的简单示例。

标签： tensorflow

解决方案

这取决于如何创建 protobuf 文件。

如果 .pb 文件是以下结果：

    # Create a builder to export the model
    builder = tf.saved_model.builder.SavedModelBuilder("export")
    # Tag the model in order to be capable of restoring it specifying the tag set
    builder.add_meta_graph_and_variables(sess, ["tag"])
    builder.save()

您必须知道该模型是如何被标记的，并使用该tf.saved_model.loader.load方法将保存的图形加载到当前的空图形中。

如果模型已被冻结，则必须手动将二进制文件加载到内存中：

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

graph = tf.get_default_graph()
tf.import_graph_def(graph_def, name="prefix")

在这两种情况下，您都必须知道输入张量的名称和要执行的节点的名称：

例如，如果您的输入张量是一个名为的占位符batch_，而您要执行的节点是dense/BiasAdd:0您必须执行的节点

    batch = graph.get_tensor_by_name('batch:0')
    prediction = restored_graph.get_tensor_by_name('dense/BiasAdd:0')

    values = sess.run(prediction, feed_dict={
        batch: your_input_batch,
    })

推荐阅读