TensorFlow2.6.0使用tf.py_function后model.fit提示数据集为空错误

运行环境 #

• Tensorflow: 2.6.0
• Ubuntu 20.04.3 LTS
• GPU: GeForce MX130
• CUDA version: 11.2

问题复现 #

使用tfio解码32位tiff时报错，改用tifffile+tf.py_function在tf.data流水线中加载数据：

TIFFReadDirectory: Warning, Unknown field with tag 42113 (0xa481) encountered.
TIFFReadDirectory: Warning, Unknown field with tag 42113 (0xa481) encountered.
memory: Sorry, can not handle images with 32-bit samples.
数据集可正常取批次可视化，但调用model.fit()时抛出错误：

ValueError: Expect x to be a non-empty array or dataset.

错误原因 #

1. tf.py_function不会自动传递返回张量的形状信息，你打印的数据集形状中通道维度为None，tf无法正确推断数据集的总长度和结构，误判数据集为空。
2. 调用model.fit()时给已经分批的tf.data数据集额外传了batch_size参数，干扰了tf对数据集结构的判断。
3. 若手动设置的steps_per_epoch、validation_steps大于数据集实际能提供的批次数量，也会触发该错误。
4. 需确认train_size + val_size + test_size 不超过总样本数，避免拆分后数据集实际为空。

修复步骤 #

1. 显式设置预处理输出的张量形状 #

在preprocess函数返回前补充形状声明：

def preprocess(img_path: str):
  # 原有预处理逻辑不变
  img = tf.image.convert_image_dtype(img, tf.float32)
  mask = tf.image.convert_image_dtype(mask, tf.uint8)

  # 新增形状设置
  img.set_shape([128, 128, 3])
  mask.set_shape([128, 128, 1])
  
  return img, mask

修改后数据集的形状信息会变为((None, 128, 128, 3), (None, 128, 128, 1))，tf可正常推断结构。

2. 修正model.fit参数 #

删掉无效的batch_size参数，确认步数计算正确，或直接删除步数参数让tf自动计算：

# 正确写法示例
model.fit(train_dataset, 
          epochs=20,
          validation_data=val_dataset,
          callbacks=callbacks)

如果需要手动指定步数，确保计算逻辑正确：

STEPS_PER_EPOCH = train_size // 10 # 批次大小为10
VALIDATION_STEPS = val_size // 10

3. （可选）显式声明数据集基数 #

如果tf仍无法识别数据集长度，可手动指定样本基数：

train_dataset = train_dataset.apply(tf.data.experimental.assert_cardinality(train_size))
val_dataset = val_dataset.apply(tf.data.experimental.assert_cardinality(val_size))

验证修复效果 #

修改后可先打印各数据集样本数确认非空：

print("训练集样本数：", len(list(train_dataset)))
print("验证集样本数：", len(list(val_dataset)))

然后重新启动训练即可。

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