Inference with TimeDistributed(Dense(1, activation="sigmoid")) in Keras 3 is a lot slower than in Keras 2 (tf.keras) Profiling shows TimeDistributed is the bottleneck.

Model: Conv1D → LSTM → Conv1DTranspose → TimeDistributed(Dense(1)).

input shape (none,none,1)

Tried: run_eagerly=False, tf.function, TensorFlow/JAX/PyTorch backends—no significant improvement.

Question

Is there a solution to fix the TimeDistributed slowdown in Keras 3 without changing the model architecture? Currently using tf_keras as a workaround.

Details

Keras Version: 3.9.0

Comment From: sonali-kumari1

Hi @Moublrs - Thanks for reporting this issue. Could you please provide a minimal code snippet to reproduce this issue?

Comment From: Moublrs

Hello @sonali-kumari1 thanks for your response here is the architecture of the model :

input_layer = Input(shape=(None, 1), name="input_layer")
x = Conv1D(32, kernel_size=kernel_size, strides=strides, padding="same",name="conv1d" )(input_layer)
x = Activation("relu", name="relu")(x)
x = Bidirectional(LSTM(32, recurrent_activation="sigmoid", return_sequences=True),name="LSTM" )(x)
x = conv1d_transpose(x, filters=32, kernel_size=kernel_size, strides=strides, padding="same")
output_layer = TimeDistributed(Dense(1, activation="sigmoid"), name="output_layer)(x)
model = Model(inputs=[input_layer], outputs=[output_layer])

the conv1d_transpose function has been taken for stackoverflow, i tested other implementations same problem, i dont think that the issue comes from ths function Thank you

Comment From: sonali-kumari1

Hi @Moublrs -

I tested the model using both keras 3 and tf.keras but I observed that keras 3 is performing slightly better than tf.keras. Computational time for keras3 is ~5.38s and tf.keras is ~6.52s for 10 inference runs. I am attaching both the gist files here:

If you are seeing different results on your end, please provide details such as the version of keras 2 you are using, or any specific hardware configurations. Thanks!

Comment From: github-actions[bot]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

Comment From: github-actions[bot]

This issue was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further.

Comment From: google-ml-butler[bot]

Are you satisfied with the resolution of your issue? Yes No

Comment From: Moublrs

here is my model the input is a time series of 50000 timestamp

def build_peak_model(config): """ Builds the peak detection model with hyperparameters from config.

Args: config (dict): Dictionary containing model hyperparameters.

Returns: keras.Model: Compiled Keras model. """

Extraction des hyperparams depuis config

# on laisse sigmoid par défaut

# Input layer
input_layer = Input(shape=(None, 1), name="input_layer")

# Downsampling
x = Conv1D(
    config["filters"],
    kernel_size=config["kernel_size"],
    strides = config["strides"],
    padding="same",
    name="conv1d"
)(input_layer)

x = Activation(config["activation_relu"], name=config["activation_relu"])(x)

# Bidirectional LSTM
x = Bidirectional(
    LSTM(
        units = config["lstm_units"],
        recurrent_activation=config["activation_sig"],
        return_sequences=True
    ),
    name="LSTM"
)(x)

# Upsampling
x = Conv1DTranspose(
    filters=config["filters"],
    kernel_size=config["kernel_size"],
    strides=config["strides"],
    padding="same",
    name="conv1d_transpose"
)(x)

# Output layer
output_layer = TimeDistributed(
    Dense(1, activation="sigmoid"), name="output_layer"
)(x)

model = Model(inputs=[input_layer], outputs=[output_layer])

return model

Keras TimeDistributed layer computational time

Extraction des hyperparams depuis config