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

today · Accepted Answer · 2019-04-04T11:50:59.173

3

You need to use a Lambda layer to perform custom operations:

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

The output of Keras layers are TF Tensors, but augmented with some additional Keras-specific attributes which is needed for constructing the model. When you directly use tf.reduce_mean, its output would be a Tensor without those additional attributes. However, when you do the same operation inside a Lambda layer, those additional attributes will be added and therefore everything would work properly.

edited Apr 04 '19 at 11:50

answered Apr 04 '19 at 11:18

today

32,602
8
95
115

this fails when saving model as h5. When loading back in, it complains about tf. – Anton Jan 13 '20 at 07:40
@Anton If it complains about `tf.reduce_mean` as not defined, use `custom_objects` argument when loading the model and pass `tf.reduce_mean`. See [this answer](https://stackoverflow.com/a/52845882/2099607). – today Jan 13 '20 at 10:27
Saw that, but a weird duct tape solution, don't you think? – Anton Jan 13 '20 at 19:58
I found that you can fix tf import by using K.mean instead of tf.reduce_mean – Anton Jan 14 '20 at 01:21
@Anton Indeed, you can use `K.mean` instead of `tf.reduce_mean` and it would work. I used `tf.reduce_mean` because the OP used it in the question. – today Jan 14 '20 at 10:04

score 0 · Answer 2 · answered Jan 13 '20 at 20:05

Objects that start with capitals inside the keras.layers module, such as Dense, LSTM.... , are layers which follow the convention of Layer(...)(input), whereas lowercase are applied directly to tensors. Useful inside of custom loss functions or custom layers (like Lambda or this example).

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

2 Answers2