GradientTape berechnet die Ausprägung in der Verlustfunktion

Question

Ich versuche, ein LSTM-Netzwerk aufzubauen , um Sätze zu klassifizieren und die Klassifizierung mit Saliency zu erklären . Dieses Netzwerk muss aus der wahren Klasse lernen y_trueund aus welchen Wörtern er nicht achten sollte Z(binäre Maske).

Dieses Papier hat uns dazu inspiriert, unsere Verlustfunktion zu entwickeln. Ich möchte, dass meine Verlustfunktion so aussieht:

Coût de classificationübersetzt in classification_lossund Coût d'explication (saillance)nach saliency_loss(was dem Gradienten der Ausgabe für die Eingabe entspricht) im folgenden Code . Ich habe versucht, dies mit einem benutzerdefinierten Modell in Keras mit Tensorflow als Backend zu implementieren:

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
        ]

Ich kann sowohl rechnen classification_lossals auch saliency_losseinen skalaren Wert erhalten. Dies funktioniert jedoch: tape.gradient(classification_loss, trainable_vars)aber das funktioniert nichttape.gradient(classification_loss + saliency_loss, trainable_vars) und wirft ValueError: No gradients provided for any variable.

Answer 1

Sie führen Berechnungen außerhalb des Bandkontexts durch (nach dem ersten gradientAufruf) und versuchen anschließend, weitere Farbverläufe zu erstellen. Das funktioniert nicht; Alle Operationen zur Unterscheidung müssen im Kontextmanager ausgeführt werden. Ich würde vorschlagen, Ihren Code wie folgt mit zwei verschachtelten Bändern umzustrukturieren:

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)

Jetzt haben wir ein Band, das für die Berechnung der Gradienten für die Eingabe für die Ausprägung verantwortlich ist. Wir haben ein weiteres Band um es herum, das diese Operationen verfolgt und später den Gradienten des Gradienten (dh den Gradienten der Ausprägung) berechnen kann. Dieses Band berechnet auch Gradienten für den Klassifizierungsverlust. Ich habe den Klassifizierungsverlust in den äußeren Bandkontext verschoben, weil das innere Band ihn nicht benötigt. Beachten Sie auch, dass selbst die Addition der beiden Verluste innerhalb des Kontex des äußeren Bandes liegt - alles muss dort geschehen, sonst geht der Berechnungsgraph verloren / unvollständig und Gradienten können nicht berechnet werden.

Answer 2

Versuchen Sie , dekorieren train_step()[e-mail]