从 TF2 中的 SavedModel 获取操作

Question

是否可以在保存的模型中获得对内部操作的引用？

我已经从 tfhub 下载了一个保存的模型，并且在检查 saved_model.pb 文件后知道我需要哪个操作。我想获得对操作的引用，以便在推理过程中记录值。

更具体地说，我想从 bert 中提取 op transformer/layer_1/dropout_5/dropout/Mul:z:0：

max_seq_length = 128
input_word_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                       name="input_word_ids")
input_mask = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                   name="input_mask")
segment_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                    name="segment_ids")

bert_layer = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/2",
                            trainable=True)

pooled_output, sequence_output = bert_layer([input_word_ids, input_mask, segment_ids])

# I think I need to add a reference to the internal op to the outputs here.
model = Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=[pooled_output, sequence_output])

Answer 1

• 详尽实用的教程 -> https://www.tensorflow.org/guide/saved_model

• Keras保存模型详细指南->
  https://www.tensorflow.org/guide/keras/save_and_serialize