ValueError: Output tensors of a Functional model must be the output of a TensorFlow `Layer` (thus ho-CodePudding

I'm having trouble with building a Functional Keras model with a customized Keras layer. The goal is to get a model whose output is that of a regular function expressed in terms of a customized Keras layer. My code is based on the ComputeSum Layer class found here.

Imports:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Layer, Input

Customized Layer (non-trainable):

class ComputeSum(Layer):

  def __init__(self, input_dim):
      super(ComputeSum, self).__init__()
      # Create a non-trainable weight.
      self.total = tf.Variable(initial_value=tf.zeros((input_dim,)),
                               trainable=False)

  def call(self, inputs):
      self.total.assign_add(tf.reduce_sum(inputs, axis=0))
      return self.total

Model:

def ComputeSumModel(tensor):
    inputs = Input(shape = tf.shape(tensor))
    outputs = ComputeSum(len(tensor))(tensor)
    
    model = keras.Model(inputs = inputs, outputs = outputs)

    return model

Testing directly on customized Layer:

these = tf.ones((4,4))

that = ComputeSum(len(these))

those = that(these)

those.numpy()

Result: array([4., 4., 4., 4.], dtype=float32)

Testing using model function:

those = ComputeSumModel(these)

Result: ValueError: Output tensors of a Functional model must be the output of a TensorFlow 'Layer' (thus holding past layer metadata). Found: <tf.Variable 'Variable:0' shape=(4,) dtype=float32, numpy=array([4., 4., 4., 4.], dtype=float32)>

What am I doing wrong here? It knows what the output of my customized layer should be, but it doesn't seem to like that output for the Functional model. Any help or suggestions would be appreciated. I'm still new to Keras and TensorFlow in general.

CodePudding user response：

You first need to create your model object before passing data into it. Also your model's input has to be fed to the ComputeSum layer. Try something like this:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Layer, Input

class ComputeSum(Layer):

  def __init__(self, input_dim):
      super(ComputeSum, self).__init__()
      # Create a non-trainable weight.
      self.total = tf.Variable(initial_value=tf.zeros((input_dim,)),
                               trainable=False)

  def call(self, inputs):
      self.total.assign_add(tf.reduce_sum(inputs, axis=0))
      return self.total

def ComputeSumModel(input_shape):
    inputs = Input(shape = input_shape)
    outputs = ComputeSum(input_shape[0])(inputs)
    
    model = keras.Model(inputs = inputs, outputs = outputs)

    return model

these = tf.ones((4,4))
that = ComputeSum(len(these))
those = that(these)
print(those)
model = ComputeSumModel((tf.shape(these)[1],))
print(model(these))

<tf.Variable 'Variable:0' shape=(4,) dtype=float32, numpy=array([4., 4., 4., 4.], dtype=float32)>
<tf.Variable 'Variable:0' shape=(4,) dtype=float32, numpy=array([4., 4., 4., 4.], dtype=float32)>