Comment répertorier toutes les opérations utilisées dans Tensorflow SavedModel?

Si tensorflow.saved_model.savej'enregistre mon modèle à l'aide de la fonction au format SavedModel, comment puis-je récupérer les Ops Tensorflow utilisées dans ce modèle par la suite. Comme le modèle peut être restauré, ces opérations sont stockées dans le graphique, ma supposition est dans le saved_model.pbfichier. Si je charge ce protobuf (donc pas le modèle entier), la partie bibliothèque du protobuf les répertorie, mais cela n'est pas documenté et étiqueté comme une fonctionnalité expérimentale pour l'instant. Les modèles créés dans Tensorflow 1.x n'auront pas cette partie.

Alors, quel est un moyen rapide et fiable pour récupérer une liste des opérations utilisées (comme MatchingFilesou WriteFile) à partir d'un modèle au format SavedModel?

En ce moment, je peux geler le tout, comme le tensorflowjs-converterfait. Comme ils vérifient également les opérations prises en charge. Cela ne fonctionne pas actuellement lorsqu'un LSTM est dans le modèle, voir ici . Y a-t-il une meilleure façon de le faire, car les opérations sont définitivement là-dedans?

Un exemple de modèle:

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

Attendu en sortie tous les Ops, contenant dans ce cas au moins:

• ReadFilecomme décrit ici
• ...

S'il saved_model.pbs'agit d'un SavedModelmessage protobuf, vous obtenez les opérations directement à partir de là. Disons que nous créons un modèle comme suit:

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

Nous pouvons maintenant trouver les opérations utilisées par ce modèle comme ceci:

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