训练DeepLabv3+时遇维度不匹配ValueError（无人机图像场景）

尝试用6000×4000分辨率的无人机JPG图像和PNG掩码训练DeepLabv3+模型时，触发ValueError：labels.shape必须与logits.shape除最后一维外保持一致，当前labels.shape=(786432,)，logits.shape=(262144, 23)。本人是深度学习新手，求技术解决方案。

import os
import cv2
import numpy as np
from glob import glob
from scipy.io import loadmat
import matplotlib.pyplot as plt

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# 数据集创建
IMAGE_SIZE = 512
BATCH_SIZE = 1
NUM_CLASSES = 23
DATA_DIR = 'C:/Users/Julian/Desktop/semantic_drone_dataset_semantics/deeplab'
NUM_TRAIN_IMAGES = 380
NUM_VAL_IMAGES = 20


train_images = sorted(glob(os.path.join(DATA_DIR, "images/*")))[:NUM_TRAIN_IMAGES]
train_masks = sorted(glob(os.path.join(DATA_DIR, "masks/*")))[:NUM_TRAIN_IMAGES]
val_images = sorted(glob(os.path.join(DATA_DIR, "images/*")))[NUM_TRAIN_IMAGES : NUM_VAL_IMAGES + NUM_TRAIN_IMAGES]
val_masks = sorted(glob(os.path.join(DATA_DIR, "masks/*")))[NUM_TRAIN_IMAGES : NUM_VAL_IMAGES + NUM_TRAIN_IMAGES]


def read_image(image_path, mask=False):
    image = tf.io.read_file(image_path)
    if mask:
        image = tf.image.decode_png(image, channels=3)
        image.set_shape([None, None, 3])
        image = tf.image.resize(images=image, size=[IMAGE_SIZE, IMAGE_SIZE])
    else:
        image = tf.image.decode_png(image, channels=3)
        image.set_shape([None, None, 3])
        image = tf.image.resize(images=image, size=[IMAGE_SIZE, IMAGE_SIZE])
        image = image / 127.5 - 1
    return image


def load_data(image_list, mask_list):
    image = read_image(image_list)
    mask = read_image(mask_list, mask=True)
    return image, mask


def data_generator(image_list, mask_list):
    dataset = tf.data.Dataset.from_tensor_slices((image_list, mask_list))
    dataset = dataset.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)
    return dataset


train_dataset = data_generator(train_images, train_masks)
val_dataset = data_generator(val_images, val_masks)

print("Train Dataset:", train_dataset)
print("Val Dataset:", val_dataset)

# 创建DeeplabV3+
def convolution_block(
    block_input,
    num_filters=256,
    kernel_size=3,
    dilation_rate=1,
    padding="same",
    use_bias=False,
):
    x = layers.Conv2D(
        num_filters,
        kernel_size=kernel_size,
        dilation_rate=dilation_rate,
        padding="same",
        use_bias=use_bias,
        kernel_initializer=keras.initializers.HeNormal(),
    )(block_input)
    x = layers.BatchNormalization()(x)
    return tf.nn.relu(x)


def DilatedSpatialPyramidPooling(dspp_input):
    dims = dspp_input.shape
    x = layers.AveragePooling2D(pool_size=(dims[-3], dims[-2]))(dspp_input)
    x = convolution_block(x, kernel_size=1, use_bias=True)
    out_pool = layers.UpSampling2D(
        size=(dims[-3] // x.shape[1], dims[-2] // x.shape[2]), interpolation="bilinear",
    )(x)

    out_1 = convolution_block(dspp_input, kernel_size=1, dilation_rate=1)
    out_6 = convolution_block(dspp_input, kernel_size=3, dilation_rate=6)
    out_12 = convolution_block(dspp_input, kernel_size=3, dilation_rate=12)
    out_18 = convolution_block(dspp_input, kernel_size=3, dilation_rate=18)

    x = layers.Concatenate(axis=-1)([out_pool, out_1, out_6, out_12, out_18])
    output = convolution_block(x, kernel_size=1)
    return output

def DeeplabV3Plus(image_size, num_classes):
    model_input = keras.Input(shape=(image_size, image_size, 3))
    resnet50 = keras.applications.ResNet50(
        weights="imagenet", include_top=False, input_tensor=model_input
    )
    x = resnet50.get_layer("conv4_block6_2_relu").output
    x = DilatedSpatialPyramidPooling(x)

    input_a = layers.UpSampling2D(
        size=(image_size // 4 // x.shape[1], image_size // 4 // x.shape[2]),
        interpolation="bilinear",
    )(x)
    input_b = resnet50.get_layer("conv2_block3_2_relu").output
    input_b = convolution_block(input_b, num_filters=48, kernel_size=1)

    x = layers.Concatenate(axis=-1)([input_a, input_b])
    x = convolution_block(x)
    x = convolution_block(x)
    x = layers.UpSampling2D(
        size=(image_size // x.shape[1], image_size // x.shape[2]),
        interpolation="bilinear",
    )(x)
    model_output = layers.Conv2D(num_classes, kernel_size=(1, 1), padding="same")(x)
    return keras.Model(inputs=model_input, outputs=model_output)


model = DeeplabV3Plus(image_size=IMAGE_SIZE, num_classes=NUM_CLASSES)
model.summary()

# 训练模型
loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss=loss,
    metrics=["accuracy"],
)

history = model.fit(train_dataset, validation_data=val_dataset, epochs=25)

plt.plot(history.history["loss"])
plt.title("Training Loss")
plt.ylabel("loss")
plt.xlabel("epoch")
plt.show()

plt.plot(history.history["accuracy"])
plt.title("Training Accuracy")
plt.ylabel("accuracy")
plt.xlabel("epoch")
plt.show()

plt.plot(history.history["val_loss"])
plt.title("Validation Loss")
plt.ylabel("val_loss")
plt.xlabel("epoch")
plt.show()

plt.plot(history.history["val_accuracy"])
plt.title("Validation Accuracy")
plt.ylabel("val_accuracy")
plt.xlabel("epoch")
plt.show()

问题根源 #

报错核心是掩码形状与模型输出不匹配：

• 模型输出logits.shape=(1, 512, 512, 23)（对应batch_size=1，512×512分辨率，23个类别）
• 你的掩码读取后是3通道RGB格式，训练时被自动展平为一维数组(786432,)，无法和模型输出的空间维度对齐

具体修复步骤 #

1. 添加RGB掩码到类别索引的转换函数

Semantic Drone Dataset的RGB值对应固定类别ID，先定义颜色映射表，再将3通道掩码转为单通道类别索引：

# 定义Semantic Drone Dataset的颜色-类别映射（需根据数据集官方文档补充完整23个类别）
COLOR_MAP = tf.constant([
    [0, 0, 0],        # 背景
    [128, 0, 0],      # 建筑
    [0, 128, 0],      # 道路
    [128, 128, 0],    # 植被
    # 其余类别颜色请参考数据集说明文件补充
])

def rgb_to_mask(rgb_image):
    # 将RGB图像转为类别索引
    rgb_image = tf.cast(rgb_image, tf.int32)
    # 计算每个像素与颜色映射的匹配
    matches = tf.equal(rgb_image[:, :, None], COLOR_MAP)
    class_mask = tf.argmax(tf.reduce_all(matches, axis=3), axis=2)
    return class_mask

2. 修改掩码读取逻辑

在read_image函数中，处理掩码时调用转换函数，输出单通道类别索引：

def read_image(image_path, mask=False):
    image = tf.io.read_file(image_path)
    if mask:
        image = tf.image.decode_png(image, channels=3)
        image.set_shape([None, None, 3])
        image = tf.image.resize(images=image, size=[IMAGE_SIZE, IMAGE_SIZE])
        # 转换为单通道类别索引
        image = rgb_to_mask(image)
        # 添加通道维度，使形状为(512,512,1)，与模型输出维度匹配
        image = tf.expand_dims(image, axis=-1)
    else:
        image = tf.image.decode_png(image, channels=3)
        image.set_shape([None, None, 3])
        image = tf.image.resize(images=image, size=[IMAGE_SIZE, IMAGE_SIZE])
        image = image / 127.5 - 1
    return image

3. 验证数据形状

可在load_data函数中添加打印，确认图像和掩码形状是否匹配：

def load_data(image_list, mask_list):
    image = read_image(image_list)
    mask = read_image(mask_list, mask=True)
    # 打印形状，确认image=(512,512,3)，mask=(512,512,1)
    tf.print("Image shape:", tf.shape(image), "Mask shape:", tf.shape(mask))
    return image, mask

额外建议 #

• 务必确保COLOR_MAP与数据集官方提供的RGB-类别对应关系完全一致，避免类别映射错误
• 若数据集提供了.mat格式的类别映射文件，可通过loadmat读取后转为TensorFlow常量，减少手动输入错误
• 训练前可视化少量图像和转换后的掩码，确认类别索引转换正确

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