训练UNet遇形状不兼容错误：Shapes (None,None)与(None,None,None,1174)

问题描述 #

搭建首个UNet模型时，训练阶段触发ValueError，提示Shapes (None, None) and (None, None, None, 1174)不兼容。尝试将损失函数改为sparse_categorical_crossentropy，并对应设置数据生成器的class_mode='sparse'后，依然出现形状匹配错误，无法广播值。

原代码 #

import tensorflow as tf
from keras import layers, models
from keras.preprocessing.image import ImageDataGenerator

def unet(input_shape=(160, 160, 3), num_classes=1174):
    inputs = tf.keras.Input(shape=input_shape)

    # Encoder (contracting path)
    conv1 = layers.Conv2D(64, 3, activation='relu', padding='same')(inputs)
    conv1 = layers.Conv2D(64, 3, activation='relu', padding='same')(conv1)
    pool1 = layers.MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = layers.Conv2D(128, 3, activation='relu', padding='same')(pool1)
    conv2 = layers.Conv2D(128, 3, activation='relu', padding='same')(conv2)
    pool2 = layers.MaxPooling2D(pool_size=(2, 2))(conv2)

    # Bottleneck
    conv3 = layers.Conv2D(256, 3, activation='relu', padding='same')(pool2)
    conv3 = layers.Conv2D(256, 3, activation='relu', padding='same')(conv3)

    # Decoder (expansive path)
    up4 = layers.UpSampling2D(size=(2, 2))(conv3)
    up4 = layers.Conv2D(128, 2, activation='relu', padding='same')(up4)
    concat4 = layers.Concatenate()([conv2, up4])
    conv4 = layers.Conv2D(128, 3, activation='relu', padding='same')(concat4)
    conv4 = layers.Conv2D(128, 3, activation='relu', padding='same')(conv4)

    up5 = layers.UpSampling2D(size=(2, 2))(conv4)
    up5 = layers.Conv2D(64, 2, activation='relu', padding='same')(up5)
    concat5 = layers.Concatenate()([conv1, up5])
    conv5 = layers.Conv2D(64, 3, activation='relu', padding='same')(concat5)
    conv5 = layers.Conv2D(64, 3, activation='relu', padding='same')(conv5)

    # Output layer
    outputs = layers.Conv2D(num_classes, 1, activation='softmax')(conv5)

    # Define the model
    model = models.Model(inputs=inputs, outputs=outputs)
    return model

train_directory = '/mnt/c/Users/user1/my_repo/data_folder/train'
test_directory = '/mnt/c/Users/user1/my_repo/data_folder/test'

# Data generators
train_datagen = ImageDataGenerator(rescale=1./255)
test_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
    train_directory,
    target_size=(160, 160),
    batch_size=8,
    class_mode='categorical'
)

test_generator = test_datagen.flow_from_directory(
    test_directory,
    target_size=(160, 160),
    batch_size=8,
    class_mode='categorical'
)

#Fit the mode;

model = unet()
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

model.fit(
    train_generator,
    steps_per_epoch=len(train_generator),
    epochs=10,
    validation_data=test_generator,
    validation_steps=len(test_generator)
)

报错信息 #

File "/home/user/.local/lib/python3.10/site-packages/keras/src/backend.py", line 5573, in categorical_crossentropy
        target.shape.assert_is_compatible_with(output.shape)

ValueError: Shapes (None, None) and (None, None, None, 1174) are incompatible

问题根源 #

你混淆了图像分类和语义分割的任务逻辑：

• UNet是语义分割模型，输出是(batch_size, 160, 160, 1174)的像素级分类结果，每个像素对应一个类别的概率分布
• 但flow_from_directory加载的是图像分类数据集，标签是单样本级别的类别标识（形状为(batch_size, 1174)或(batch_size,)），和UNet的输出形状完全不匹配，导致损失函数计算时形状不兼容

修复方案 #

1. 调整数据集结构 #

语义分割需要图像与掩码一一配对的数据集，标准结构如下：

train/
  images/
    img_001.jpg
    img_002.jpg
    ...
  masks/
    mask_001.png  # 与img_001对应，每个像素值代表该位置的类别索引（0~1173）
    mask_002.png
    ...
test/
  images/
  masks/

2. 自定义数据加载逻辑 #

替换flow_from_directory，改用能同时加载图像和对应掩码的代码：

import os
import numpy as np
from PIL import Image
import tensorflow as tf

def load_segmentation_dataset(img_dir, mask_dir, target_size=(160,160), num_classes=1174):
    # 匹配图像和掩码文件（确保文件名对应）
    img_paths = sorted([os.path.join(img_dir, f) for f in os.listdir(img_dir) if f.lower().endswith(('.png','.jpg','.jpeg'))])
    mask_paths = sorted([os.path.join(mask_dir, f) for f in os.listdir(mask_dir) if f.lower().endswith(('.png','.jpg','.jpeg'))])

    images = []
    masks = []
    for img_path, mask_path in zip(img_paths, mask_paths):
        # 加载并预处理图像
        img = Image.open(img_path).resize(target_size)
        img = np.array(img, dtype=np.float32) / 255.0
        
        # 加载并预处理掩码（用最近邻插值避免类别混淆）
        mask = Image.open(mask_path).resize(target_size, Image.Resampling.NEAREST)
        mask = np.array(mask, dtype=np.int32)
        
        # 可选：转换为one-hot编码（如果使用categorical_crossentropy）
        # mask = tf.keras.utils.to_categorical(mask, num_classes=num_classes)
        
        images.append(img)
        masks.append(mask)
    
    return np.array(images), np.array(masks)

# 加载训练和测试数据
train_imgs, train_masks = load_segmentation_dataset(
    '/mnt/c/Users/user1/my_repo/data_folder/train/images',
    '/mnt/c/Users/user1/my_repo/data_folder/train/masks'
)
test_imgs, test_masks = load_segmentation_dataset(
    '/mnt/c/Users/user1/my_repo/data_folder/test/images',
    '/mnt/c/Users/user1/my_repo/data_folder/test/masks'
)

3. 调整模型编译与训练 #

根据掩码的格式选择对应的损失函数：

• 如果掩码是整数索引格式（单通道，像素值0~1173），用sparse_categorical_crossentropy
• 如果掩码是one-hot编码格式（多通道，每个通道对应一类），用categorical_crossentropy

示例代码：

model = unet()

# 整数掩码的情况推荐用sparse_categorical_crossentropy（节省内存）
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=[tf.keras.metrics.MeanIoU(num_classes=1174)]  # 语义分割专用指标，比普通accuracy更合适
)

# 训练模型
model.fit(
    train_imgs, train_masks,
    batch_size=8,
    epochs=10,
    validation_data=(test_imgs, test_masks)
)

4. 额外注意事项 #

• 确保掩码的像素值范围是0到num_classes-1，没有超出范围的无效值，否则会触发索引错误
• 1174类属于极多类别场景，建议分批加载数据（比如用tensorflow.data.Dataset实现流式加载），避免一次性加载全部数据占用过多内存
• 语义分割任务中，普通的accuracy指标参考价值低，优先使用MeanIoU、Dice系数等专用评估指标

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