GradientTape calcule la prominencia en la función de pérdida
Estoy tratando de construir una red LSTM para clasificar oraciones y proporcionar una explicación para la clasificación usando prominencia . Esta red debe aprender de la clase verdadera y_trueasí como de qué palabras no debe prestar atención Z(máscara binaria).
Este documento nos inspiró a idear nuestra función de pérdida. Así es como me gustaría que se viera mi función de pérdida:
Coût de classificationse traduce en classification_lossy Coût d'explication (saillance)a saliency_loss(que es el mismo que el gradiente de salida wrt la entrada) en el código de abajo . Intenté implementar esto con un modelo personalizado en Keras, con Tensorflow como 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
]
Me las arreglo para calcular classification_losstan bien como saliency_lossy obtengo un valor escalar. Sin embargo, esto funciona: tape.gradient(classification_loss, trainable_vars)pero esto no funcionatape.gradient(classification_loss + saliency_loss, trainable_vars) y arroja ValueError: No gradients provided for any variable.
Respuestas
Está haciendo cálculos fuera del contexto de la cinta (después de la primera gradientllamada) y luego está tratando de tomar más gradientes después. Esto no funciona; todas las operaciones para diferenciar deben ocurrir dentro del administrador de contexto. Sugeriría reestructurar su código de la siguiente manera, usando dos cintas anidadas:
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)
Ahora tenemos una cinta responsable de calcular los gradientes con la entrada para la prominencia. Tenemos otra cinta a su alrededor que rastrea esas operaciones y luego podemos calcular el gradiente del gradiente (es decir, gradiente de la prominencia). Esta cinta también calcula gradientes para la pérdida de clasificación. Moví la pérdida de clasificación en el contexto de la cinta exterior porque la cinta interior no la necesita. Tenga en cuenta también cómo incluso la suma de las dos pérdidas está dentro del contexto de la cinta exterior: todo tiene que suceder allí, de lo contrario, el gráfico de cálculo se pierde / está incompleto y los gradientes no se pueden calcular.
Intenta decorar train_step()con@tf.function