时间戳

首页 > 解决方案 > 从 TF2 中的 SavedModel 获取操作

python - 从 TF2 中的 SavedModel 获取操作

问题描述

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

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

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

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

标签： pythontensorflowkerastensorflow2.0

解决方案

推荐阅读