如何在Tensorflow SavedModel中列出所有使用的操作?


10

如果我使用tensorflow.saved_model.saveSavedModel格式的函数保存模型,那么以后如何检索该模型中使用的Tensorflow Ops。由于可以还原模型,因此这些操作存储在图中,我的猜测在saved_model.pb文件中。如果我加载此protobuf(而不是整个模型),则protobuf的库部分会列出这些,但目前尚未记录和标记为实验功能。在Tensorflow 1.x中创建的模型将没有此部分。

那么,从SavedModel格式的模型中检索使用过的操作列表(如MatchingFilesWriteFile)的快速可靠的方法是什么?

现在,我可以像冻结一样冻结整个事情tensorflowjs-converter。因为他们还会检查受支持的操作。当模型中包含LSTM时,这当前不起作用,请参见此处。因为Ops肯定在那里,有没有更好的方法呢?

示例模型:

class FileReader(tf.Module):

@tf.function(input_signature=[tf.TensorSpec(name='filename', shape=[None], dtype=tf.string)])
def read_disk(self, file_name):
    input_scalar = tf.reshape(file_name, [])
    output = tf.io.read_file(input_scalar)
    return tf.stack([output], name='content')

file_reader = FileReader()

tf.saved_model.save(file_reader, 'file_reader')

预期输出所有操作,在这种情况下至少包含:


1
很难确切地说出您想要什么saved_model.pb,它是tf.GraphDef还是SavedModelprotobuf信息?如果您有tf.GraphDef通话gd,您可以使用来获得使用过的操作列表sorted(set(n.op for n in gd.node))。如果您有加载的模型,则可以执行sorted(set(op.type for op in tf.get_default_graph().get_operations()))。如果是SavedModel,则可以从中获取tf.GraphDef(例如saved_model.meta_graphs[0].graph_def)。
jdehesa

我想从存储的SavedModel中检索操作。所以确实,您正在描述的最后一个选项。saved_model您上一个示例中的变量是什么?结果tf.saved_model.load('/path/to/model')或加载saved_model.pb文件的protobuf的。
sampers

Answers:


1

如果saved_model.pbSavedModelprotobuf消息,那么您可以直接从那里获取操作。假设我们创建一个模型,如下所示:

import tensorflow as tf

class FileReader(tf.Module):
    @tf.function(input_signature=[tf.TensorSpec(name='filename', shape=[None], dtype=tf.string)])
    def read_disk(self, file_name):
        input_scalar = tf.reshape(file_name, [])
        output = tf.io.read_file(input_scalar)
        return tf.stack([output], name='content')

file_reader = FileReader()
tf.saved_model.save(file_reader, 'tmp')

现在,我们可以找到该模型使用的操作,如下所示:

from tensorflow.core.protobuf.saved_model_pb2 import SavedModel

saved_model = SavedModel()
with open('tmp/saved_model.pb', 'rb') as f:
    saved_model.ParseFromString(f.read())
model_op_names = set()
# Iterate over every metagraph in case there is more than one
for meta_graph in saved_model.meta_graphs:
    # Add operations in the graph definition
    model_op_names.update(node.op for node in meta_graph.graph_def.node)
    # Go through the functions in the graph definition
    for func in meta_graph.graph_def.library.function:
        # Add operations in each function
        model_op_names.update(node.op for node in func.node_def)
# Convert to list, sorted if you want
model_op_names = sorted(model_op_names)
print(*model_op_names, sep='\n')
# Const
# Identity
# MergeV2Checkpoints
# NoOp
# Pack
# PartitionedCall
# Placeholder
# ReadFile
# Reshape
# RestoreV2
# SaveV2
# ShardedFilename
# StatefulPartitionedCall
# StringJoin

我尝试了类似的方法,但是不幸的是,这并没有达到我的预期:说我有一个可以执行此input_scalar = tf.reshape(file_name, []) output = tf.io.read_file(input_scalar) return tf.stack([output], name='content')操作的模型:然后,这里列出的ReadFile Op 在那里,但是没有打印出来。
凌晨

1
@sampers我已经用您建议的示例编辑了答案。我确实ReadFile在输出中得到了操作。在您的实际情况下,该操作是否可能不在保存的模型的输入和输出之间?在那种情况下,我认为它可能会被修剪。
jdehesa

实际上,使用给定的模型就可以了。不幸的是,对于在tf2中制作的模块,事实并非如此。如果我创建一个tf.Module 1个功能与file_name参数@tf.function标注,包含我在以前的评论中列出的电话,它提供了以下列表:Const, NoOp, PartitionedCall, Placeholder, StatefulPartitionedCall
sampers

增加了一个模型,我的问题
sampers

@sampers我已经更新了答案。我以前使用的是TF 1.x,我不熟悉TF 2.x中对图形定义对象的更改,我认为答案现在涵盖了已保存模型中的所有内容。我认为与您编写的Python函数相对应的操作在其中saved_model.meta_graphs[0].graph_def.library.function[0](该node_def函数对象内的集合)。
jdehesa
By using our site, you acknowledge that you have read and understand our Cookie Policy and Privacy Policy.
Licensed under cc by-sa 3.0 with attribution required.