如何在 Tensorflow SavedModel 中列出所有使用过的操作? [英] How to list all used operations in Tensorflow SavedModel?

问题描述

如果我使用 tensorflow.saved_model.save 函数以 SavedModel 格式保存我的模型，那么我之后如何检索此模型中使用了哪些 Tensorflow Ops.由于模型可以恢复，所以这些操作都存储在图中，我猜是在saved_model.pb文件中.如果我加载这个 protobuf(所以不是整个模型)，protobuf 的库部分会列出这些，但目前没有记录并标记为实验功能.在 Tensorflow 1.x 中创建的模型将没有这部分.

If I save my model using the tensorflow.saved_model.save function in SavedModel format, how can I retrieve which Tensorflow Ops are used in this model afterwards. As the model can be restored, these operations are stored in the graph, my guess is in the saved_model.pb file. If I load this protobuf (so not the entire model) the library part of the protobuf lists these, but this is not documented and tagged as an experimental feature for now. Models created in Tensorflow 1.x won't have this part.

那么，从 SavedModel 格式的模型中检索已使用操作列表(如 MatchingFiles 或 WriteFile)的快速可靠方法是什么?

So what is a fast and reliable way to retrieve a list of used Operations (Like MatchingFiles or WriteFile) from a model in SavedModel format?

现在我可以冻结整个内容，比如 tensorflowjs-converter 可以.因为他们还检查支持的操作.当 LSTM 在模型中时，这当前不起作用，请参阅此处.有没有更好的方法来做到这一点，因为运营人员肯定在那里?

Right now I can freeze the entire thing, like tensorflowjs-converter does. As they also check for supported Operations. This currently does not work when an LSTM is in the model, see here. Is there a better way to do this, as the Ops are definitely in there?

示例模型:

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')

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

Expected in output all Ops, containing in this case at least:

• ReadFile 如此处所述
• ...

推荐答案

如果 saved_model.pb 是 SavedModel protobuf 消息，那么您可以直接从那里获得操作.假设我们创建一个模型如下:

If saved_model.pb is a SavedModel protobuf message, then you get the operations directly from there. Let's say we create a model as follows:

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')

我们现在可以像这样找到该模型使用的操作:

We can now find the operations used by that model like this:

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

这篇关于如何在 Tensorflow SavedModel 中列出所有使用过的操作?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！