如何解决TensorFlow 2.8中pix2pix模型的MeanIoU混淆矩阵错误？

MeanIoU指标要求输入的预测值必须是非负的类别索引整数，但pix2pix模型的输出通常是未经过argmax的logits，或是经过tanh激活后范围在[-1,1]的张量，这直接触发了断言检测到负数的错误。修改num_classes无法解决问题，因为核心矛盾是输入格式不符合要求，而非类别数量设置错误。

方案1：自定义适配的MeanIoU指标 #

编写自定义指标类，在计算前先将模型输出转换为符合要求的类别索引，不影响原有训练逻辑：

import tensorflow as tf
from tensorflow.keras.losses import DiceLoss

class CustomMeanIoU(tf.keras.metrics.MeanIoU):
    def update_state(self, y_true, y_pred, sample_weight=None):
        # 将模型输出的多通道logits转换为类别索引（最后一维为类别维度）
        y_pred = tf.argmax(y_pred, axis=-1)
        # 统一标签与预测值的整数类型
        y_true = tf.cast(y_true, tf.int32)
        y_pred = tf.cast(y_pred, tf.int32)
        # 调用父类的指标更新逻辑
        super().update_state(y_true, y_pred, sample_weight)

编译模型时替换原指标：

myTransformer.compile(optimizer='adam',
                      loss=DiceLoss(),
                      metrics=[CustomMeanIoU(num_classes=100)])

方案2：针对单通道输出的特殊处理 #

如果你的pix2pix模型是单通道输出（比如二分类、灰度级分割），且输出经过tanh激活到[-1,1]范围，可以这样映射转换：

class CustomMeanIoU(tf.keras.metrics.MeanIoU):
    def update_state(self, y_true, y_pred, sample_weight=None):
        # 将[-1,1]映射到[0, num_classes-1]区间并取整为类别索引
        y_pred = (y_pred + 1) / 2 * (self.num_classes - 1)
        y_pred = tf.cast(tf.round(y_pred), tf.int32)
        y_true = tf.cast(y_true, tf.int32)
        super().update_state(y_true, y_pred, sample_weight)

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