GradientTape oblicza istotność w funkcji straty

Question

Próbuję zbudować sieć LSTM, aby klasyfikować zdania i wyjaśniać klasyfikację za pomocą istotności . Ta sieć musi uczyć się od prawdziwej klasy, y_truea także od jakich słów nie powinien zwracać uwagi Z(maska ​​binarna).

Ten artykuł zainspirował nas do wymyślenia naszej funkcji straty. Oto, jak chciałbym, aby wyglądała moja funkcja utraty:

Coût de classificationtłumaczy na classification_lossi Coût d'explication (saillance)na saliency_loss(co jest tym samym, co gradient wyjścia na wejściu) w poniższym kodzie . Próbowałem to zaimplementować z niestandardowym modelem w Keras, z Tensorflow jako zapleczem:

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
        ]

Udaje mi się classification_lossrównie dobrze obliczyć saliency_lossi otrzymuję wartość skalarną. Jednak to działa: tape.gradient(classification_loss, trainable_vars)ale to nie działatape.gradient(classification_loss + saliency_loss, trainable_vars) i wyrzuca ValueError: No gradients provided for any variable.

Answer 1

Wykonujesz obliczenia poza kontekstem taśmy (po pierwszym gradientwywołaniu), a następnie próbujesz pobrać więcej gradientów. To nie działa; wszystkie operacje rozróżniania muszą być wykonywane wewnątrz menedżera kontekstu. Sugerowałbym zmianę struktury kodu w następujący sposób, używając dwóch zagnieżdżonych taśm:

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)

Teraz mamy jedną taśmę odpowiedzialną za obliczanie gradientów na wejściu do istotności. Mamy wokół niego inną taśmę, która śledzi te operacje i może później obliczyć gradient gradientu (tj. Gradient istotności). Taśma oblicza również gradienty dla utraty klasyfikacji. Przeniosłem utratę klasyfikacji w kontekście taśmy zewnętrznej, ponieważ taśma wewnętrzna tego nie potrzebuje. Zwróć również uwagę, że nawet dodanie dwóch strat znajduje się w kontekście zewnętrznej taśmy - wszystko musi się tam wydarzyć, w przeciwnym razie wykres obliczeniowy zostanie utracony / niekompletny i nie można obliczyć gradientów.

Answer 2

Spróbuj ozdobić train_step()[e-mail]