python - Can't save Keras model when using reduce_mean in the model?

Question

The Keras code snippet reads:

second_input = inputs_d['second_input']
selected = embedding_layer(second_input)
item_average = tf.reduce_mean(selected, axis=1, keepdims=True)
second_input_encoded = tf.keras.layers.Reshape((3,))(item_average)

If I change second_input as from shape(5,) to shape (1,) and get rid of reduce_mean the code runs just fine.

The error message reads:

/site-packages/tensorflow/python/util/serialization.py", line 69, in get_json_type raise TypeError('Not JSON Serializable:', obj) TypeError: ('Not JSON Serializable:', b"\n\x04Mean\x12\x04Mean\x1a'embedding_1/embedding_lookup/Identity_2\x1a\x16Mean/reduction_indices*\x07\n\x01T\x12\x020\x01*\n\n\x04Tidx\x12\x020\x03*\x0f\n\tkeep_dims\x12\x02(\x01")

标签： pythontensorflowkeras

Answer 1

您需要使用一个Lambda层来执行自定义操作：

item_average = tf.keras.layers.Lambda(lambda x: tf.reduce_mean(x, axis=1, keepdims=True))(selected)

Keras 层的输出是 TF 张量，但增加了一些额外的 Keras 特定属性，这些属性是构建模型所需的。当您直接使用tf.reduce_mean时，它的输出将是一个没有这些附加属性的张量。但是，当您在Lambda层内执行相同操作时，将添加这些附加属性，因此一切都会正常工作。