Paso de entrenamiento personalizado de TensorFlow con diferentes funciones de pérdida

Question

Antecedentes

Según la documentación de TensorFlow , se puede realizar un paso de entrenamiento personalizado con lo siguiente

# Fake sample data for testing
x_batch_train = tf.zeros([32, 3, 1], dtype="float32")
y_batch_train = tf.zeros([32], dtype="float32")

loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
with tf.GradientTape() as tape:
    logits = model(x_batch_train, training=True)
    loss_value = loss_fn(y_batch_train, logits)

grads = tape.gradient(loss_value, model.trainable_weights)
optimizer.apply_gradients(zip(grads, model.trainable_weights))

Pero si quiero usar una función de pérdida diferente como la entropía cruzada categórica, necesitaría argmax los logits creados en la cinta de degradado:

loss_fn = tf.keras.lossees.get("categorical_crossentropy")
with tf.GradientTape() as tape:
    logits = model(x_batch_train, training=True)
    prediction = tf.cast(tf.argmax(logits, axis=-1), y_batch_train.dtype)
    loss_value = loss_fn(y_batch_train, prediction)

grads = tape.gradient(loss_value, model.trainable_weights)
optimizer.apply_gradients(zip(grads, model.trainable_weights))

---

Problema

El problema con esto es que la tf.argmaxfunción no es diferenciable, por lo que TensorFlow no podría calcular los gradientes y obtendría el error:

ValueError: No gradients provided for any variable: [...]

---

Mi pregunta: sin cambiar la función de pérdida, ¿cómo podría hacer que el segundo ejemplo funcione?

Answer 1

categorical_crossentropy espera que sus etiquetas tengan una codificación en caliente, por lo que debe asegurarse de eso primero. Luego pase directamente el resultado de su modelo, esta salida debe ser una probabilidad por categoría más información ->https://www.tensorflow.org/api_docs/python/tf/keras/losses/CategoricalCrossentropy#standalone_usage