Creating BLEU loss method on tensorflow gives "No gradient provided"

user_8576878 0
Viewed 302

I need to build a custom loss method based on BLEU. I'm passing my LabelEncoder in the constructor to reverse labels and predictions and calculate the bleu distance.

Here is my Loss class

class CIMCodeSuccessiveLoss(Loss):

    def __init__(self, labelEncoder: LabelEncoder):
        super().__init__()
        self.le = labelEncoder

    def bleu_score(self, true_label, pred_label):
        cim_true_label = self.le.inverse_transform(true_label.numpy())
        cim_pred_label = self.le.inverse_transform(pred_label.numpy())
        bleu_scores = [sentence_bleu(list(one_true_label),
                                     list(one_pred_label),
                                     weights=(0.5, 0.25, 0.125, 0.125)) for one_true_label, one_pred_label in
                       zip(cim_true_label, cim_pred_label)]
        return np.float32(bleu_scores)

    def call(self, y_true, y_pred):
        labeled_y_pred = tf.cast(tf.argmax(y_pred, axis=-1), tf.int32)
        bleu = tf.py_function(self.bleu_score, (tf.reshape(y_true, [-1]), labeled_y_pred), tf.float32)
        return tf.reduce_sum(tf.square(1 - bleu))

The bleu_score method is calculating the correct scores and returns a NumPy array. when I try to return the squared sum, I get this error

raise ValueError(f"No gradients provided for any variable: {variable}.

I'm also providing the model:

inputs = tf.keras.Input(shape=(1,), dtype=tf.string)
x = vectorize_layer(inputs)
x = Embedding(vocab_size, embedding_dim, name="embedding")(x)
x = LSTM(units=32, name="lstm")(x)
outputs = Dense(classes_number, name="classification")(x)

model = tf.keras.Model(inputs=inputs, outputs=outputs, name="first_cim_classifier")

model.summary()


# we add early stopping for our model.
early_stopping = EarlyStopping(monitor='loss', patience=2)

model.compile(
    loss=CIMCodeSuccessiveLoss(le),
    optimizer=tf.keras.optimizers.Adam(),
    metrics=["accuracy", "crossentropy"],
    run_eagerly=True)

trained_model = model.fit(np.array(x_train), np.array(y_train), batch_size=64, epochs=10,
                          validation_data=(np.array(x_val), np.array(y_val)),
                          callbacks=[early_stopping])

Any help is appreciated. Thanks in advance.

tensorflow deep-learning python
2 Answers

To calculate the loss function, you use the method 'tf.argmax(y_pred, axis=-1)',argmax is not differentiable and the automatic differentiation to calculate the gradients is not possible, you have to remove this method, for example (depending on your data) you can change the output layer to softmax and labels to one_hot.

user_52355280

The issue is, the argmax function is not a differentiable, which is problematic when including it in a loss function:

labeled_y_pred = tf.cast(tf.argmax(y_pred, axis=-1), tf.int32)

One way to workaround this is to use a differentiable approximation of the argmax function, similar to the smooth maximum function:

enter image description here

As β approaches infinity, this will approach the the true maximum. For your purposes, β=10 or β=100 should accomplish your goals.

In Tensorflow, this could be accomplished as follows:

def differentiable_argmax_approx(x, beta=10, axis=None):
    return tf.reduce_sum(tf.cumsum(tf.ones_like(x)) * tf.exp(beta * x) / tf.reduce_sum(tf.exp(beta * x), axis=axis), axis=axis) - 1

Then changing the original line to:

labeled_y_pred = tf.cast(differentiable_argmax_approx(y_pred, axis=-1), tf.int32)

We can verify the functionality with a simple test case:

    beta = 10
    x = np.array([1, 2, 3, 10, 4, 5], dtype=np.float)
    y = differentiable_argmax_approx(x, beta)
    assert x.argmax() == y

One caveat to this approach: if the maximum value is not unique along the axis that we're applying the function to, the result will be the arithmetic mean of the indices. Providing another test case to illustrate:

    beta = 10
    x = np.array([1, 2, 10, 3, 10], dtype=np.float)
    y = differentiable_argmax_approx(x, beta)
    assert y == 3

The result is 3 here, because we have two occurrences of the maximum value (10): one at index 2, and the other at index 4. In contrast, the regular argmax function returns the first index of the maximum argument.

Another improvement would be moving more computation into Tensorflow functions. To start, instead of using sklearn's LabelEncoder, to apply a mapping in the loss function, you could use a tf.lookup.StaticHashTable to accomplish the same objective with the Tensorflow API. To convert from a LabelEncoder to a tf.lookup.StaticHashTable, you can use the following function:

def convert_label_encoder_to_static_hash_table(le: LabelEncoder,
                                         default_value: int = -1) -> tf.lookup.StaticHashTable:
    static_hash_table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(
            tf.convert_to_tensor(le.classes_),
            tf.convert_to_tensor(le.transform(le.classes_))), default_value=default_value)
    return static_hash_table

Or, for your purposes, since you're applying the inverse mapping (to go from integers -> string), you may want to swap the key and the values:

def convert_label_encoder_to_static_hash_table(le: LabelEncoder,
                                         default_value: int = "") -> tf.lookup.StaticHashTable:
    static_hash_table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(
            tf.convert_to_tensor(le.transform(le.classes_)),
            tf.convert_to_tensor(le.classes_))), default_value=default_value)
    return static_hash_table

and, in the initializer:

    def __init__(self, labelEncoder: LabelEncoder):
        super().__init__()
        self.table = convert_label_encoder_to_static_hash_table(labelEncoder)

By operating on tf.Tensor objects, you can utilize tf.map_fn instead of using a for-loop and converting to a numpy array/lists - your loss function would become:

def bleu_score(self, true_label, pred_label):
        cim_true_label = self.table[true_label]
        cim_pred_label = self.table[pred_label]
        bleu_scores = tf.map_fn(lambda x: sentence_bleu([str(x[0])], [str(x[1])], weights=(0.5, 0.25, 0.125, 0.125)),
                            elems=tf.stack([(ground_truth, pred) for ground_truth, pred in
                                            zip(cim_pred_label, cim_true_label)],
                                           dtype=(tf.string, tf.string),
                                     fn_output_signature=tf.int32))
        return bleu_scores

This should also mitigate the need to call tf.py_func in the loss computation, since the bleu_score function is now entirely Tensorflow operations instead of calling native Python functions.

user_130736770
Related