需要初始化的tf.keras自定义层无法使用tf.contrib.saved_model.save_keras_model保存

时间:2019-01-01 22:32:59

标签: python keras

我正在尝试实现一个自定义查找层,该层将字符串转换为ints。我需要将模型保存为张量流服务格式。该模型需要一个需要初始化的查找表。在自定义层构建定义中,使用tf.keras.backend.get_session() session初始化表。这样可以很好地训练,但是在使用tf.contrib.saved_model.save_keras_model保存时会引发以下错误:

ValueError: Cannot use the given session to execute operation: the operation's graph is different from the session's graph.

以下代码重现该错误:     将numpy导入为np     将tensorflow导入为tf

class LookupLayer(tf.keras.layers.Layer):
  def __init__(self, mapping=[''], num_oov_buckets=0, default_value=-1, **kwargs):
    self.mapping=mapping
    self.num_oov_buckets=num_oov_buckets
    self.default_value=default_value
    super(LookupLayer, self).__init__(**kwargs)

  def build(self, input_shape):
    self.index_table = tf.contrib.lookup.index_table_from_tensor(
                    mapping=self.mapping,
                    num_oov_buckets=self.num_oov_buckets,
                    default_value=self.default_value,
                )
    self.index_table.init.run(session=tf.keras.backend.get_session())
    super(LookupLayer, self).build(input_shape)

  def call(self, input):
    return self.index_table.lookup(input)

  def compute_output_shape(self, input_shape):
    return input_shape

input = tf.keras.layers.Input(shape=(1,), dtype="string")
lookup_output = LookupLayer(mapping=['test'], num_oov_buckets=1)(input)
emb_layer = tf.keras.layers.Embedding(2, 1)(lookup_output)
x = tf.keras.layers.Flatten()(emb_layer)
x = tf.keras.layers.Dense(100, activation='relu')(x)
out = tf.keras.layers.Dense(1, activation='sigmoid')(x)
model = tf.keras.models.Model(inputs=input, outputs=out)
model.compile(optimizer=tf.train.AdamOptimizer(),
              loss='binary_crossentropy')

X={'input_1':np.array(['test', 'oov'])}
y=[0,1]
model.fit(X,y)

tf.contrib.saved_model.save_keras_model(model, './saved_model_test/', custom_objects={'LookupLayer': LookupLayer})

如何与自定义tf.keras层进行正确的会话?还是有更好的方法来做到这一点?

注意:我需要字符串-> int查找在图中。我无法在单独的预处理步骤中使用它,因为我需要它在那儿进行服务。

1 个答案:

答案 0 :(得分:1)

我能够通过使用simple_save而不是save_keras_model将模型保存到pb文件中进行服务:

"verbose": true

注意:请确保使用legacy_init_op = tf.tables_initializer(),而不要像How to keep lookup tables initialized for prediction (and not just training)?中的答案那样使用legacy_init_op = tf.saved_model.main_op.main_op()。否则,似乎所有的权重都会重置,并且您的模型无法投放。

这不能解决我为之发表这篇文章的问题(save_keras_model不起作用),但可以解决我的用例。