UNet医学图像分割训练出现input image must be of non-zero size错误原因是什么

问题核心原因&修复方案 #

错误根因

1. 自定义数据生成器输出格式不符合模型要求：你当前的生成器每次返回单张图片，形状为(512,512,1)，缺少TensorFlow要求的第一维batch维度，传入上采样层时会被判定为尺寸异常。
2. model.fit参数缺失：使用无限循环的自定义生成器训练时，必须指定steps_per_epoch和validation_steps参数，告知模型每轮训练/验证的迭代次数。
3. 损失函数不匹配：你最后一层是sigmoid激活的单通道输出，属于二分类分割任务，应该使用BinaryCrossentropy损失，当前使用的SparseCategoricalCrossentropy是多分类任务专用，会引发计算异常。
4. 额外风险：如果你的DICOM文件总数量不足，总切片数小于1200，会导致lungs_test = lungs[1200:]得到空数组，也会触发非零尺寸报错。

修复步骤

1. 先运行以下代码确认数据加载正常：

print(lungs.shape, seg.shape)
# 正常输出格式应为 (总样本数, 512, 512, 1) (总样本数, 512, 512, 1)
# 如果总样本数小于1200，需要调整训练集/测试集的拆分比例，或者补充DICOM数据

2. 修改自定义数据生成器，补充batch维度：

def get_gen(data_x: list, data_y: list, batch_size=8):
  assert len(data_x) == len(data_y), "数据长度不匹配"
  while True:
    for i in range(0, len(data_x), batch_size):
      batch_x = data_x[i:i+batch_size]
      batch_y = data_y[i:i+batch_size]
      yield np.array(batch_x), np.array(batch_y)

# 初始化生成器
batch_size = 8
lungs_tr_gen = get_gen(lungs_train, seg_train, batch_size=batch_size)
lungs_ts_gen = get_gen(lungs_test, seg_test, batch_size=batch_size)

3. 修改模型编译的损失函数：

# 替换原模型定义里的compile部分
model.compile(optimizer='adam',
              loss=tf.keras.losses.BinaryCrossentropy(),
              metrics=['accuracy'])

4. 修改model.fit调用，补充步数参数：

steps_per_epoch = len(lungs_train) // batch_size
validation_steps = len(lungs_test) // batch_size

model.fit(lungs_tr_gen, 
          epochs=10, 
          steps_per_epoch=steps_per_epoch,
          validation_data=lungs_ts_gen,
          validation_steps=validation_steps)

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