GradientTape calcule la saillance dans la fonction de perte

Question

J'essaie de créer un réseau LSTM pour classer les phrases et expliquer la classification en utilisant la saillance . Ce réseau doit apprendre de la vraie classe y_trueainsi que des mots auxquels il ne doit pas prêter attention Z(masque binaire).

Cet article nous a inspiré pour concevoir notre fonction de perte. Voici à quoi j'aimerais que ma fonction de perte ressemble:

Coût de classificationse traduit par classification_losset Coût d'explication (saillance)en saliency_loss(qui est le même que le gradient de sortie par rapport à l'entrée) dans le code ci-dessous . J'ai essayé de l'implémenter avec un modèle personnalisé dans Keras, avec Tensorflow comme backend:

loss_tracker = metrics.Mean(name="loss")
classification_loss_tracker = metrics.Mean(name="classification_loss")
saliency_loss_tracker = metrics.Mean(name="saliency_loss")
accuracy_tracker = metrics.CategoricalAccuracy(name="accuracy")

class CustomSequentialModel(Sequential):
        
    def _train_test_step(self, data, training):
        # Unpack the data
        X = data[0]["X"]
        Z = data[0]["Z"] # binary mask (1 for important words)
        y_true = data[1]
        
        # gradient tape requires "float32" instead of "int32"
        # X.shape = (None, MAX_SEQUENCE_LENGTH, EMBEDDING_DIM)
        X = tf.cast(X, tf.float32)

        # Persitent=True because we call the `gradient` more than once
        with GradientTape(persistent=True) as tape:
            # The tape will record everything that happens to X
            # for automatic differentiation later on (used to compute saliency)
            tape.watch(X)
            # Forward pass
            y_pred = self(X, training=training) 
            
            # (1) Compute the classification_loss
            classification_loss = K.mean(
                categorical_crossentropy(y_true, y_pred)
            )
 
            # (2) Compute the saliency loss
            # (2.1) Compute the gradient of output wrt the maximum probability
            log_prediction_proba = K.log(K.max(y_pred))
            
        # (2.2) Compute the gradient of the output wrt the input
        # saliency.shape is (None, MAX_SEQUENCE_LENGTH, None)
        # why isn't it (None, MAX_SEQUENCE_LENGTH, EMBEDDING_DIM) ?!
        saliency = tape.gradient(log_prediction_proba, X)
        # (2.3) Sum along the embedding dimension
        saliency = K.sum(saliency, axis=2)
        # (2.4) Sum with the binary mask
        saliency_loss = K.sum(K.square(saliency)*(1-Z))
        # =>  ValueError: No gradients provided for any variable
        loss = classification_loss + saliency_loss 
        
        trainable_vars = self.trainable_variables
        # ValueError caused by the '+ saliency_loss'
        gradients = tape.gradient(loss, trainable_vars) 
        del tape # garbage collection
        
        if training:
            # Update weights
            self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        
        # Update metrics
        saliency_loss_tracker.update_state(saliency_loss)
        classification_loss_tracker.update_state(classification_loss)
        loss_tracker.update_state(loss)
        accuracy_tracker.update_state(y_true, y_pred)
        
        # Return a dict mapping metric names to current value
        return {m.name: m.result() for m in self.metrics}
    
    def train_step(self, data):
        return self._train_test_step(data, True)
    
    def test_step(self, data):
        return self._train_test_step(data, False)
    
    @property
    def metrics(self):
        return [
            loss_tracker,
            classification_loss_tracker,
            saliency_loss_tracker,
            accuracy_tracker
        ]

J'arrive à calculer classification_lossaussi bien saliency_losset j'obtiens une valeur scalaire. Cependant, cela fonctionne: tape.gradient(classification_loss, trainable_vars)mais cela ne fonctionne pastape.gradient(classification_loss + saliency_loss, trainable_vars) et jette ValueError: No gradients provided for any variable.

Answer 1

Vous effectuez des calculs en dehors du contexte de la bande (après le premier gradientappel) et essayez ensuite de prendre plus de dégradés par la suite. Cela ne fonctionne pas; toutes les opérations de différenciation doivent avoir lieu dans le gestionnaire de contexte. Je suggérerais de restructurer votre code comme suit, en utilisant deux bandes imbriquées:

with GradientTape() as loss_tape:
    with GradientTape() as saliency_tape:
        # The tape will record everything that happens to X
        # for automatic differentiation later on (used to compute saliency)
        saliency_tape.watch(X)
        # Forward pass
        y_pred = self(X, training=training) 
        
        # (2) Compute the saliency loss
        # (2.1) Compute the gradient of output wrt the maximum probability
        log_prediction_proba = K.log(K.max(y_pred))
        
    # (2.2) Compute the gradient of the output wrt the input
    # saliency.shape is (None, MAX_SEQUENCE_LENGTH, None)
    # why isn't it (None, MAX_SEQUENCE_LENGTH, EMBEDDING_DIM) ?!
    saliency = saliency_tape.gradient(log_prediction_proba, X)
    # (2.3) Sum along the embedding dimension
    saliency = K.sum(saliency, axis=2)
    # (2.4) Sum with the binary mask
    saliency_loss = K.sum(K.square(saliency)*(1-Z))

    # (1) Compute the classification_loss
    classification_loss = K.mean(
        categorical_crossentropy(y_true, y_pred)
    )

    loss = classification_loss + saliency_loss 
    
trainable_vars = self.trainable_variables
gradients = loss_tape.gradient(loss, trainable_vars)

Maintenant, nous avons une bande responsable du calcul des gradients par rapport à l'entrée pour la saillance. Nous avons une autre bande autour d'elle qui suit ces opérations et peut ensuite calculer le gradient du gradient (c'est-à-dire le gradient de la saillance). Cette bande calcule également les gradients pour la perte de classification. J'ai déplacé la perte de classification dans le contexte de la bande externe car la bande interne n'en a pas besoin. Notez également comment même l'addition des deux pertes se trouve à l'intérieur des contextes de la bande externe - tout doit s'y passer, sinon le graphe de calcul est perdu / incomplet et les gradients ne peuvent pas être calculés.

Answer 2

Essayez de décorer train_step()[e-mail]