keras.layers.Layer.call method fails when building keras model with functional API

[pmot01]

I am trying to write a custom convolution layer by subclassing keras.layers.Layer. As long as I use it as part of a Sequential model and run it eagerly, it seems to work. But trying to build the model with the functional API it causes a failure within my call method:
ValueError: Tried to convert 'y' to a tensor and failed. Error: None values not supported.

What´s going on when building the model using the functional API?

I wonder how to write a call method that works within te functional API. Here´s my code, I simplyfied the code in "call" so as to focus on the failing part.

import keras
import numpy as np
from keras import layers
import tensorflow as tf
import typing

# Using the functional API it fails during model construction. If building a Sequential model it fails
# with the same error at the beginning of the training (in fit).
# Running functions eagerly changes the picture: Using the functional API still fails with the same error
# at the same location, but the Sequential model is executed without an exception!
use_functional_api = True
use_eager_mode = True


tf.config.run_functions_eagerly(use_eager_mode)


ActivationFunction = typing.Callable[[typing.Any], typing.Any] | str

INPUT_SAMPLES_INDEX = 0
INPUT_HEIGHT_INDEX = 1
INPUT_WIDTH_INDEX = 2
INPUT_CHANNELS_INDEX = 3


class MyConv2D(keras.layers.Layer):
def __init__(self, *, num_filters: int, kernel_size: int, activation_function: ActivationFunction, **kwargs):
super().__init__(**kwargs) # pass on name, dtype, ... parameters

tf.print("MyConv2D.init", num_filters, kernel_size, activation_function, kwargs)
self.num_filters = num_filters
self.kernel_size = kernel_size
self.activationFunction = keras.activations.get(activation_function) \
if isinstance(activation_function, str) else activation_function

# Values calculated in build when input shape is known
self.kernel_shape = None # The same kernel used for all images
self.activation_map_shape = None # The resulting activation map for one image
# Weights
self.w = None # Kernel weights for all filters
self.b = None # Bias weights for all filters

# In the build method is where the weights are constructed
def build(self, input_shape: tf.TensorShape):
# the kernel has the given spatial dimension and the number of channels from the input and the desired depth
# of filters
tf.print("MyConv2D.build", input_shape)
self.kernel_shape = (self.kernel_size, self.kernel_size, input_shape[INPUT_CHANNELS_INDEX], self.num_filters)
result_height = self.compute_result_dimension(input_shape[INPUT_HEIGHT_INDEX])
result_width = self.compute_result_dimension(input_shape[INPUT_WIDTH_INDEX])
self.activation_map_shape = (result_height, result_width, self.num_filters)

self.w = self.add_weight(shape=self.kernel_shape, initializer="random_normal", trainable=True)
self.b = self.add_weight(shape=(self.num_filters, 1), initializer="random_normal", trainable=True)

def compute_result_dimension(self, input_dimension):
return input_dimension - self.kernel_size + 1

# def __call__(self, *args, **kwargs):
# tf.print("MyDense.__call__")
# return super().__call__(*args, **kwargs)

def call(self, input_samples: tf.Tensor, *args, **kwargs):
tf.print("MyConv2D.call", input_samples.shape)

num_input_samples = input_samples.shape[INPUT_SAMPLES_INDEX]
output_samples = []
i = tf.constant(0, dtype=tf.int32)
while tf.less(i, num_input_samples):
output_samples.append(self.transform_sample(input_samples[i]))
i += 1
return tf.stack(output_samples)

def transform_sample(self, sample):
# Just to simplify debugging
return tf.random.normal(self.activation_map_shape, mean=0.0, stddev=1.0)

def compute_output_shape(self, input_shape):
tf.print("MyConv2D.compute_output_shape", input_shape)
return (input_shape[0],) + self.activation_map_shape

def get_config(self):
config = super().get_config()
config.update({
'num_filters': self.num_filters,
'kernel_size': self.kernel_size,
'activation_function': self.activation_function.__name__
})
return config


original_max_grayscale_value = 255
image_width_in_pixels = 28
image_height_in_pixels = 28
image_channels = 1 # grayscale, with uint8 values between 0 and 255, need to be scaled to float32 between 0 and 1


def reshape_images(images: np.ndarray) -> np.ndarray:
# The convolutional layers are designed to process image data, so we don´t need to flatten the images. But we
# need to add the channels layer, which is 1 here (that´s why it was possible to leave it out in the first place)
resulting_images = images.reshape((len(images), image_width_in_pixels, image_width_in_pixels, image_channels))
return resulting_images.astype("float32") / original_max_grayscale_value


print("##### Load and prepare data #####")
(train_images, train_labels), (test_images, test_labels) = keras.datasets.mnist.load_data()
train_images = reshape_images(train_images)
test_images = reshape_images(test_images)
# Use only a few for research
actual_train_images = train_images[:5]
actual_train_labels = train_labels[:5]
actual_test_images = test_images[:5]
actual_test_labels = test_labels[:5]

print("##### Build the model #####")
common_kernel_size = 3
common_pool_size = 2

myLayer = MyConv2D(num_filters=32, kernel_size=common_kernel_size, activation_function="relu", name="Conv_1")
result = myLayer(actual_train_images)
print(f"Resulting type: {result.__class__} and shape: {result.shape}")

if use_functional_api:
inputs = keras.Input(shape=(image_height_in_pixels, image_width_in_pixels, image_channels), name="Input")
layer_stack = (MyConv2D(num_filters=32, kernel_size=common_kernel_size, activation_function="relu", name="Conv_1")
(inputs))
layer_stack = layers.Flatten(name="Flatten")(layer_stack)
outputs = layers.Dense(10, activation="softmax", name="Classifier")(layer_stack)
model = keras.Model(inputs=inputs, outputs=outputs)
model.summary()
else:
model = keras.Sequential(
[
myLayer,
layers.Flatten(name="Flatten"),
layers.Dense(10, activation="softmax", name="Classifier")
]
)

model.compile(
optimizer="rmsprop",
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
run_eagerly=use_eager_mode
)

print("##### Train the model #####")
num_epochs = 2

history = model.fit(
actual_train_images, actual_train_labels,
epochs=num_epochs,
batch_size=64
)
print(f"training history: {history.history}")

[Abhas Kumar Sinha]

Hello,

This seems like a bug, can you post this Code to https://www.github.com/keras-team/keras/issues/ in a proper format so that we can look into it?

Thank you.

[pmot01]

After upgrading to TensorFlow version 2.16.1 it works now in eager mode