基于Keras和TensorFlow的U-Net图像分割：logits与labels维度不匹配问题

错误详情 #

Node: 'sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits'
logits and labels must have the same first dimension, got logits shape [2097152,2] and labels shape [2000000]

背景 #

基于Keras官方《Image segmentation with a U-Net-like architecture》教程做图像分割，原教程用160×160图像时运行正常。仅修改了图像路径读取代码（如下），替换为自己的250×250图像与掩码数据，未改动网络架构，触发上述错误。

修改后的路径读取代码 #

target_dir = "IA_training_data_final/Toy_mask/"
img_size = (250, 250)
class_list = os.listdir(input_dir)
num_classes = len(class_list)
target_classes = list(range(num_classes))
batch_size = 32
input_img_number = 0
target_img_number = 0
input_img_paths = list()
target_img_paths = list()
val_percent = 0.10

for subdir, dirs, files in os.walk(input_dir):
    for file in files:
        input_img_number += 1
        input_path = os.path.join(subdir,file)
        input_img_paths.append(input_path)
input_img_paths = sorted(input_img_paths)

for subdir, dirs, files in os.walk(target_dir):
    for file in files:
        target_img_number += 1
        target_path = os.path.join(subdir,file)
        target_img_paths.append(target_path)
target_img_paths = sorted(target_img_paths)

print("Number of samples:", input_img_number)
print("Number of masks:", target_img_number)

问题原因与解决方法 #

核心原因 #

原教程的U-Net架构依赖2的整数次幂尺寸，而250不是2的整数次幂，经过网络多次池化（尺寸减半）和上采样（尺寸加倍）后，输出的特征图尺寸会自动调整为最近的2的整数次幂（256×256），和输入的250×250不匹配。

对应错误里的数值：

• 掩码总像素数：32（batch_size）×250×250=2000000，对应labels的shape
• 网络输出总像素数：32×256×256=2097152，对应logits的shape

两者维度不一致，导致损失函数计算报错。

解决办法 #

1. 调整输入尺寸为2的整数次幂：把图像和掩码统一调整为256×256（或128×128、512×512等），这样网络的下采样和上采样过程能完美匹配输入输出尺寸，无需改动网络架构。
2. 修改网络最后一层：在输出层前添加Cropping2D(cropping=((3,3),(3,3)))层，将256×256的输出裁剪为250×250，和输入尺寸对齐。
3. 检查数据加载环节：确认数据生成器在加载图像和掩码时，没有自动缩放或裁剪尺寸，确保输入图像和掩码严格保持250×250。

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