Làm thế nào để liệt kê tất cả các hoạt động được sử dụng trong Tensorflow SavingModel?

Nếu tôi lưu mô hình của mình bằng cách sử dụng tensorflow.saved_model.savechức năng ở định dạng SavingModel, làm thế nào tôi có thể truy xuất Ops Tensorflow nào được sử dụng trong mô hình này sau đó. Khi mô hình có thể được khôi phục, các hoạt động này được lưu trữ trong biểu đồ, tôi đoán là trong saved_model.pbtệp. Nếu tôi tải protobuf này (không phải toàn bộ mô hình) thì phần thư viện của protobuf liệt kê những thứ này, nhưng hiện tại nó không được ghi lại và được gắn thẻ như một tính năng thử nghiệm. Các mô hình được tạo trong Tensorflow 1.x sẽ không có phần này.

Vậy đâu là cách nhanh chóng và đáng tin cậy để lấy danh sách các Hoạt động đã sử dụng (Thích MatchingFileshoặc WriteFile) từ một mô hình ở định dạng SavingModel?

Ngay bây giờ tôi có thể đóng băng toàn bộ, như thế tensorflowjs-converter. Vì họ cũng kiểm tra các hoạt động được hỗ trợ. Điều này hiện không hoạt động khi một LSTM trong mô hình, xem tại đây . Có cách nào tốt hơn để làm điều này không, vì Ops chắc chắn ở trong đó?

Một mô hình ví dụ:

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

Dự kiến ​​trong đầu ra tất cả các Ops, chứa trong trường hợp này ít nhất là:

• ReadFilenhư được mô tả ở đây
• ...

Nếu saved_model.pblà một SavedModelthông điệp protobuf, thì bạn có được các hoạt động trực tiếp từ đó. Giả sử chúng ta tạo một mô hình như sau:

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

Bây giờ chúng ta có thể tìm thấy các hoạt động được sử dụng bởi mô hình đó như thế này:

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