How to average a layer's output in tensorflow?

Viewed 234

This is a toy model I am trying to implement with tensorflow. The input is a set (10) of real number pairs. And the underlying function I want to approximate is formula. The implemented model should look like this:

model architecture

I also need to mention that "Hidden Layer" is the same layer (same parameters) for all X_i.

What I implemented so far:

import tensorflow as tf
import numpy as np

def tf_model():
    # Define the inputs
    inputs = tf.keras.Input(shape=[10, 2])

    # Define common hidden layer
    hidden_layer = tf.keras.layers.Dense(64, activation="relu")(inputs)

    # Propagate and average
    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(hidden_layer)
    outputs = tf.keras.layers.Average()(outputs)

    return tf.keras.Model(inputs=inputs, outputs=output)

X = np.random.rand(1000,10,2) * 100
y = 1 / (1 + X[...,0]**2 + X[...,1]**4)
y = np.average(y, axis=1)

model = tf_model()
model.fit(X, y)

What I get from running this:

Traceback (most recent call last):
File "model_test.py", line 21, in <module>
    model = tf_model()
File "model_test.py", line 13, in tf_model
    outputs = tf.keras.layers.Average()(outputs)
File "/home/redbull/.local/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
File "/home/redbull/.local/lib/python3.8/site-packages/keras/layers/merge.py", line 88, in build
    raise ValueError(
ValueError: A merge layer should be called on a list of inputs. Received: input_shape=(None, 10, 1) (not a list of shapes)

I think the issue is that tf.keras.layers.Average() only works with a list of inputs, but not a tf layer/ tensor.

Since tf.keras.layers.Average() does not seem to be suitable in this scenario, how can I implement the wished functionality?

1 Answers

You can use tf.reduce_mean as below.

outputs = tf.reduce_mean(outputs, axis=1)
Related