如何调整TensorFlow官方教程中的图像分割模型以实现二值掩码输出？

我完全懂你现在的困扰——跟着官方Unet-style的分割教程改二值掩码模型时，踩了好几个莫名其妙的坑：设2类直接报错，设1类要么输出空掩码，要么训练时epoch越多效果越差，明明准确率还飙到0.99对吧？咱们一步步把模型调整成适配二值（0=非掩码，1=掩码）任务的样子：

一、先把模型输出层与损失函数配对应 #

二值分割的核心是单通道输出+Sigmoid激活，这是你之前踩的第一个坑：

• 不要设OUTPUT_CLASSES=2：二值分割不需要2个通道，单通道输出通过Sigmoid激活后，每个像素输出0-1的概率，刚好对应“非掩码/掩码”的概率，后续用0.5阈值就能转成0/1的掩码
• 你已经给输出层加了activation='sigmoid'，这步是对的，保持OUTPUT_CLASSES=1就好
• 损失函数用binary_crossentropy是正确的，但要注意from_logits=False（因为已经加了Sigmoid，不是原始logits输出）

修改后的输出层代码和模型初始化：

OUTPUT_CLASSES = 1

def unet_model(output_channels:int):
    inputs = layers.Input(shape=[128, 128, 3])
    skips = down_stack(inputs)
    x = skips[-1]
    skips = reversed(skips[:-1])

    for up, skip in zip(up_stack, skips):
        x = up(x)
        concat = layers.Concatenate()
        x = concat([x, skip])

    # 输出层保持Sigmoid激活
    last = tf.keras.layers.Conv2DTranspose(
        filters=output_channels, kernel_size=3, strides=2,
        padding='same', activation='sigmoid')  #64x64 -> 128x128
    x = last(x)

    return Model(inputs=inputs, outputs=x)

model = unet_model(output_channels=OUTPUT_CLASSES)
model.compile(optimizer='adam',
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=False),
              metrics=['accuracy'])

二、解决“训练越久效果越差但准确率虚高”的核心问题 #

你遇到的“epoch越多效果越差但准确率很高”，90%是类别不平衡+训练不稳定导致的：

1. 准确率虚高的本质：如果数据集里非掩码区域占比极大（比如99%都是背景），模型只要全输出0，就能拿到极高的准确率，但完全没学到掩码的特征
   • 解决方法：用加权二值交叉熵，给掩码区域更高的权重，比如如果掩码区域占比10%，就给类别1（掩码）设9倍的权重：

     # 假设掩码区域占训练集的10%，权重设为9
     class_weights = {0: 1, 1: 9}
     # 训练时传入class_weight参数
     model.fit(..., class_weight=class_weights)
     

2. 训练不稳定（越训越差）：你加了BatchNormalization有改善，这是因为BN能稳定训练过程，避免梯度震荡/消失，建议把BN加到更多位置：

   for up, skip in zip(up_stack, skips):
       x = up(x)
       x = tf.keras.layers.BatchNormalization()(x)  # 上采样后加BN
       concat = layers.Concatenate()
       x = concat([x, skip])
       x = tf.keras.layers.BatchNormalization()(x)  # 拼接特征后也加BN
   

3. 学习率太高：Adam默认的0.001学习率对于迁移学习的模型来说可能太高了，导致训练后期参数震荡，建议调低到0.0001：

   optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)
   model.compile(optimizer=optimizer, loss='binary_crossentropy', metrics=['accuracy'])
   

三、其他能提升效果的小调整 #

1. 不要只看准确率，换用分割专用指标：准确率对于分割任务参考性极低，建议加上IoU（交并比）或Dice系数，能真实反映掩码的分割质量：

   def iou_metric(y_true, y_pred):
       # 先把预测值转成0/1的二值掩码
       y_pred = tf.cast(y_pred > 0.5, tf.float32)
       intersection = tf.reduce_sum(y_true * y_pred)
       union = tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) - intersection
       # 加1e-7避免除0
       return intersection / (union + 1e-7)
   
   model.compile(optimizer=optimizer,
                 loss='binary_crossentropy',
                 metrics=['accuracy', iou_metric])
   

2. 微调预训练模型的顶层：你之前把down_stack.trainable=False，完全冻结了MobileNetV2的权重，如果你的数据集足够大（比如几百张带掩码的图），可以解冻模型的顶层几层，让预训练特征更好地适配你的任务：

   down_stack.trainable = True
   # 只解冻MobileNetV2的后20层，前面的底层特征是通用的，不用动
   for layer in base_model.layers[:-20]:
       layer.trainable = False
   

3. 预测后处理：模型输出的是0-1的概率图，记得用阈值转成0/1的二值掩码：

   import numpy as np
   
   # 预测概率图
   pred_prob = model.predict(input_image)
   # 用0.5阈值转成二值掩码
   pred_mask = (pred_prob > 0.5).astype(np.uint8)
   

备注：内容来源于stack exchange，提问作者Viacheslav Rud'