TensorFlow重写train_step()后self.metrics缺失编译指标的解决方法

你重写train_step()后，编译时指定的指标无法在fit日志和history中显示，核心原因是没有利用Keras内置的self.compiled_metrics来统一管理指标的更新与结果返回。原生Keras的训练流程会自动通过compiled_metrics处理所有编译时配置的指标，包括自定义指标和loss关联统计。

修改后的train_step实现 #

将你的train_step()方法替换为以下代码，即可完全复现原生行为：

def train_step(self, data):
    x, y = data

    with tf.GradientTape() as tape:
        ŷ = self(x, training=True)
        loss_value = self.compiled_loss(y, ŷ, regularization_losses=self.losses)

    # 计算梯度并更新权重
    gradients = tape.gradient(loss_value, self.trainable_variables)
    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

    # 更新所有编译时指定的指标（包括分类准确率）
    self.compiled_metrics.update_state(y, ŷ)

    # 构建返回结果：包含loss和所有指标
    result = {m.name: m.result() for m in self.compiled_metrics}
    result["loss"] = loss_value
    return result

关键修改说明 #

• 自动更新指标：使用self.compiled_metrics.update_state(y, ŷ)替代手动遍历self.metrics，该方法会自动处理所有编译时通过model.compile(metrics=...)指定的指标，无需逐个判断更新。
• 统一结果返回：先通过self.compiled_metrics.result()获取所有指标的计算结果，再手动加入loss值（因为loss不属于compiled_metrics的管理范畴），确保返回的字典包含原生train_step会输出的所有字段。

修改后的运行效果 #

执行fit后，日志会同时显示训练集的loss和categorical_accuracy：

1875/1875 [==============================] - 12s 5ms/step - loss: 0.2838 - categorical_accuracy: 0.9172 - val_loss: 0.1499 - val_categorical_accuracy: 0.9525

同时history.history字典中会包含loss、categorical_accuracy、val_loss、val_categorical_accuracy这些键，完全匹配原生Keras的行为。

内容的提问来源于stack exchange，提问作者u2gilles