背景

根据TensorFlow文档，可以使用以下方式执行自定义训练步骤

# 用于测试的伪样本数据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))

但如果我想使用不同的损失函数，比如类别交叉熵，我需要在梯度带中对生成的logits进行argmax操作：

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))

问题

这样做的问题在于tf.argmax函数不可微分，因此TensorFlow无法计算梯度，你会得到以下错误：

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

我的问题： 在不更改损失函数的情况下，如何使第二个示例工作？

回答：

categorical_crossentropy期望你的标签是一热编码的，所以你首先需要确保这一点。然后直接传递模型的结果，输出应该是每个类别的概率，更多信息 -> https://www.tensorflow.org/api_docs/python/tf/keras/losses/CategoricalCrossentropy#standalone_usage