Как перечислить все используемые операции в Tensorflow SavedModel?

10

Если я tensorflow.saved_model.saveсохраню свою модель, используя функцию в формате SavedModel, как я могу узнать, какие операции Tensorflow используются в этой модели впоследствии. Поскольку модель может быть восстановлена, эти операции хранятся в графике, я думаю, в saved_model.pbфайле. Если я загружаю этот protobuf (то есть не всю модель), библиотечная часть protobuf перечисляет их, но это пока не задокументировано и не помечено как экспериментальная функция. Модели, созданные в Tensorflow 1.x, не будут иметь этой части.

Итак, что такое быстрый и надежный способ получения списка используемых операций (например, MatchingFilesили WriteFile) из модели в формате SavedModel?

Прямо сейчас я могу заморозить всю вещь, как 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')

Ожидаемые в выводе все Ops, содержащие в этом случае как минимум:

sampers
источник
1
Трудно точно сказать, что вы хотите, что saved_model.pbэто tf.GraphDef, или SavedModelсообщение protobuf? Если у вас есть 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
Я хочу получить ops из сохраненной SavedModel. Так что, действительно, последний вариант, который вы описываете. Какая saved_modelпеременная в вашем последнем примере? Результат tf.saved_model.load('/path/to/model')или загрузка протобуфа файла сохраненного_моделя.pb.
пробники

Ответы:

1

Если saved_model.pbэто SavedModelсообщение protobuf, то вы получаете операции непосредственно оттуда. Допустим, мы создаем модель следующим образом:

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
jdehesa
источник
Я попробовал что-то вроде этого, но, к сожалению, это не то, что я ожидаю: скажем, у меня есть модель, которая делает это: input_scalar = tf.reshape(file_name, []) output = tf.io.read_file(input_scalar) return tf.stack([output], name='content')тогда Op ReadFile, как указано здесь, находится там, но не печатается.
пробники
1
@ sampers Я отредактировал ответ с примером, как вы предлагаете. Я получаю ReadFileоперацию на выходе. Возможно ли, что в вашем реальном случае эта операция не находится между входом и выходом сохраненной модели? В этом случае я думаю, что это может быть сокращено.
jdehesa
Действительно с данной моделью это работает. К сожалению, для модуля, сделанного в TF2, это не так. Если я создаю модуль tf.Module с 1 с аннотацией file_nameаргумента @tf.function, содержащей вызовы, которые я перечислил в моем предыдущем комментарии, он дает следующий список:Const, NoOp, PartitionedCall, Placeholder, StatefulPartitionedCall
семплеры
добавил модель к моему вопросу
пробники
@sampers Я обновил свой ответ. Я использовал TF 1.x раньше, я не был знаком с изменениями в объектах определения графиков в TF 2.x, я думаю, что ответ теперь охватывает все в сохраненной модели. Я думаю, что операции, соответствующие написанной вами функции Python, находятся в saved_model.meta_graphs[0].graph_def.library.function[0]( node_defколлекции внутри этого объекта функции).
Джедеса