How use properly tensorflow functions within the model

user_13418745 0
Viewed 400

I'm using the functional API of TensorFlow 2 and tensorflow.keras.layers to build the model.

I have an input tensor (in_1) with shape [batch_size, length, dim] and I would like to compute the mean along the length dimension and obtain an output tensor (out_1) with shape [batch_size, dim].

Which of this should I use to do it? (all these options works, in terms of output shape and training)

        out_1 = Lambda(lambda x: tf.math.reduce_mean(x, axis=1))(in_1)

        out_1 = Lambda(lambda x: tf.keras.backend.mean(x, axis=1))(in_1)

        out_1 = tf.math.reduce_mean(in_1, axis=1)

This last one automatically creates a TensorFlowOpLayer, is this something that should be avoided?

Are there other ways to do this?

What's the difference between tf.math.reduce_mean and tf.keras.backend.mean, which should I use?

I know that custom functions should be called inside the Lambda layer, but is it true also for TensorFlow functions such as tf.math.reduce_mean which can process the tensor in "one fell swoop"? How should I call them if I need to specify a parameter (e.g. axis)?

tensor keras tensorflow model python
1 Answers

First, for the difference between tf.keras.backend.mean and tf.math.reduce_mean: There is none. You can check the source code for the keras backend version, which simply uses reduce_mean (from math_ops, but internally that's the same one that's exposed in tf.math). IMHO this is a bit of a failure in the TF re-design where they incorporated Keras: Keras is now contained in TF, but Keras also uses TF in the "backend", so you basically have every operation twice: Once the TF version, and once the Keras version which, after all, also just uses the TF version.

Anyway, for the difference between using Lambda or not: It also doesn't (really) matter. Here is a minimal example:

inp = tf.keras.Input((10,))
layer = tf.reduce_mean(inp, axis=-1)

model = tf.keras.Model(inp, layer)

print(model.layers)

gives the output

[<tensorflow.python.keras.engine.input_layer.InputLayer at 0x7f1a651500b8>,
 <tensorflow.python.keras.engine.base_layer.TensorFlowOpLayer at 0x7f1a9912d8d0>]

We can see that the reduce_mean operation was automatically converted to a TensorFlowOpLayer. Now, this may be technically different from a Lambda layer, but I doubt that this makes any practical difference. I suppose this would not work for a Sequential model, where you need to supply a list of layers, so there Lambda would likely be needed.

user_93931020
Related