TensorFlow中二类图像分割CNN模型输入维度不匹配问题

问题背景 #

将原59类图像分割模型改为二分类时，出现训练输入形状不匹配错误，同时对损失函数和输出层设置存在疑惑。

一、数据处理与规格 #

掩码二值化代码 #

for path in tqdm(msk_paths):
    file = tf.io.read_file(path)
    mask = tf.image.decode_png(file,channels=1,dtype=tf.uint8)
    mask = tf.round(mask/255)
    masks.append(mask)

数据形状与类型 #

X_train.element_spec, y_train.element_spec, X_val.element_spec, y_val.element_spec

# 输出结果
(TensorSpec(shape=(128, 128, 3), dtype=tf.float32, name=None),
 TensorSpec(shape=(128, 128, 1), dtype=tf.uint8, name=None),
 TensorSpec(shape=(128, 128, 3), dtype=tf.float32, name=None),
 TensorSpec(shape=(128, 128, 1), dtype=tf.uint8, name=None))

二、模型定义（基于DenseNet121的U-Net） #

base = keras.applications.DenseNet121(input_shape=[128,128,3], 
                                      include_top=False, 
                                      weights='imagenet')

# 用于跳跃连接的特征层
skip_names = ['conv1/relu', # 尺寸64*64
             'pool2_relu',  # 尺寸32*32
             'pool3_relu',  # 尺寸16*16
             'pool4_relu',  # 尺寸8*8
             'relu'        # 尺寸4*4
             ]

skip_outputs = [base.get_layer(name).output for name in skip_names]
downstack = keras.Model(inputs=base.input, outputs=skip_outputs)
downstack.trainable = False

!pip install -q git+https://github.com/tensorflow/examples.git --quiet
from tensorflow_examples.models.pix2pix import pix2pix

# 上采样层
upstack = [pix2pix.upsample(512,3),
          pix2pix.upsample(256,3),
          pix2pix.upsample(128,3),
          pix2pix.upsample(64,3)]

# 输入层定义
inputs = keras.layers.Input(shape=[128,128,3])

# 下采样
down = downstack(inputs)
out = down[-1]

# 处理跳跃连接
skips = reversed(down[:-1])
for up, skip in zip(upstack,skips):
    out = up(out)
    out = keras.layers.Concatenate()([out,skip])
    
# 最终输出层（1通道对应二分类）
out = keras.layers.Conv2DTranspose(1, 3,
                                  strides=2,
                                  padding='same',
                                  )(out)

unet = keras.Model(inputs=inputs, outputs=out)

# 模型编译
def Compile_Model():
    unet.compile(loss=keras.losses.SparseCrossentropy(from_logits=True),
            optimizer=keras.optimizers.RMSprop(learning_rate=0.001),
            metrics=['accuracy']) 
Compile_Model()

三、训练报错信息 #

训练代码 #

hist_1 = unet.fit(train,
               validation_data=val,
               steps_per_epoch=STEPS_PER_EPOCH,
               validation_steps=VALIDATION_STEPS,
               epochs=20,
               verbose=2)

错误提示 #

ValueError: Input 0 of layer "model_10" is incompatible with the layer: expected shape=(None, 128, 128, 3), found shape=(128, 128, 3)

解决方案 #

1. 修复输入形状不匹配问题 #

报错原因是训练数据集缺少批量维度：TensorFlow模型默认期望输入形状为(batch_size, height, width, channels)，但当前数据是单样本形状(128,128,3)。

修复方式：

• 如果使用tf.data.Dataset构建数据集，添加batch()方法：

train = train.batch(32)  # 批量大小可根据显存调整，比如16/32/64
val = val.batch(32)

2. 损失函数与输出层选择（针对二分类场景） #

两种方案均可行，根据掩码格式选择：

方案A：保留当前SparseCrossentropy配置

你的掩码是单通道整数类型（0/1），SparseCrossentropy(from_logits=True)完全适配：

• 无需给输出层添加激活函数，损失会自动处理logits到概率的转换
• 输出层设置为1通道是正确的

方案B：改用BinaryCrossentropy

若要切换为该损失函数，需修改两处：

• 输出层添加sigmoid激活：

out = keras.layers.Conv2DTranspose(1, 3,
                                  strides=2,
                                  padding='same',
                                  activation='sigmoid')(out)

• 编译时更换损失函数：

unet.compile(loss=keras.losses.BinaryCrossentropy(),
            optimizer=keras.optimizers.RMSprop(learning_rate=0.001),
            metrics=['accuracy'])

注意：二分类场景不需要用softmax激活，softmax适用于多分类（类别互斥）的场景。

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