背景
根据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
