TypeError: Expected keras.losses.Loss, found function - tensorflow

Question

I want to build a TFF model for speech recognition systems. For this, I use the CNN-GRU model architecture with a CTC loss function. but I got error when I wanted to build_federated_averaging_process and think it's about the ctc_loss function but I cant fix it.
part of my code is:
def CTCLoss(y_true, y_pred):
# Compute the training-time loss value
batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64")
input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64")
label_length = tf.cast(tf.shape(y_true)[1], dtype="int64")
input_length = input_length * tf.ones(shape=(batch_len, 1), dtype="int64")
label_length = label_length * tf.ones(shape=(batch_len, 1), dtype="int64")
loss = keras.backend.ctc_batch_cost(y_true, y_pred, input_length, label_length)
return loss
def create_compiled_keras_model():
"""Model similar to DeepSpeech2."""
# Model's input
input_spectrogram = layers.Input((None, fft_length // 2 + 1), name="input")
# Expand the dimension to use 2D CNN.
x = layers.Reshape((-1, fft_length // 2 + 1 , 1), name="expand_dim")(input_spectrogram)
# Convolution layer 1
x = layers.Conv2D(
filters=32,
kernel_size=[11, 41],
strides=[2, 2],
padding="same",
use_bias=False,
name="conv_1",
)(x)
x = layers.BatchNormalization(name="conv_1_bn")(x)
x = layers.ReLU(name="conv_1_relu")(x)
# Convolution layer 2
x = layers.Conv2D(
filters=32,
kernel_size=[11, 21],
strides=[1, 2],
padding="same",
use_bias=False,
name="conv_2",
)(x)
x = layers.BatchNormalization(name="conv_2_bn")(x)
x = layers.ReLU(name="conv_2_relu")(x)
# Reshape the resulted volume to feed the RNNs layers
x = layers.Reshape((-1, x.shape[-2] * x.shape[-1]))(x)
# RNN layers
for i in range(1, 2 + 1):
recurrent = layers.GRU(
units=128,
activation="tanh",
recurrent_activation="sigmoid",
use_bias=True,
return_sequences=True,
reset_after=True,
name=f"gru_{i}",
)
x = layers.Bidirectional(
recurrent, name=f"bidirectional_{i}", merge_mode="concat"
)(x)
if i < 2:
x = layers.Dropout(rate=0.5)(x)
# Dense layer
x = layers.Dense(units=128 * 2, name="dense_1")(x)
x = layers.ReLU(name="dense_1_relu")(x)
x = layers.Dropout(rate=0.5)(x)
# Classification layer
output = layers.Dense(units= output_dim + 1, activation="softmax")(x)
# Model
model = keras.Model(input_spectrogram, output, name="DeepSpeech_2")
return model
def model_fn():
# We _must_ create a new model here, and _not_ capture it from an external
# scope. TFF will call this within different graph contexts.
keras_model = create_compiled_keras_model()
return tff.learning.from_keras_model(
keras_model,
input_spec=layers.Input((None, fft_length // 2 + 1)),
loss=CTCLoss)
and I got error in this step :
iterative_process = tff.learning.build_federated_averaging_process(
model_fn,
client_optimizer_fn=lambda:keras.optimizers.Adam(learning_rate=1e-4))
TypeError: Expected keras.losses.Loss, found function.
how do I fix it?

Answer 1

class Customloss(tf.keras.losses.Loss):
def __init__(self):
super().__init__()
#tf.function
def CTCLoss(self, y_true, y_pred):
...#
return loss
try to use tf.keras.losses.Loss for custom loss in tff. It will work.