hamid ebrahimi
hamid ebrahimi

Reputation: 51

TypeError: Expected keras.losses.Loss, found function

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?

Upvotes: 1

Views: 211

Answers (1)

Amandeep Singh
Amandeep Singh

Reputation: 41

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.

Upvotes: 0

Related Questions