TensorFlow自定义数据集训练报错：Attempted to pad to a smaller size

问题背景 #

使用TensorFlow基于自定义数据集训练模型，数据集包含2个输入、2个输出：

• 输入1：shape(320, 180, 1)的单通道图像
• 输入2：shape()的布尔值
• 输出1：shape(2,)的float32类型数据
• 输出2：shape()的整数类型数据

训练直接失败，报错信息如下：

2 root error(s) found.
(0) DATA_LOSS: Attempted to pad to a smaller size than the input element.
[[{{node IteratorGetNext}}]]
[[IteratorGetNext/_6]]
(1) DATA_LOSS: Attempted to pad to a smaller size than the input element.
[[{{node IteratorGetNext}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_1655]

已确认所有图像均为320宽、180高的8位单通道图像，无损坏数据，但使用tensorflow_datasets构建并加载数据集时，即使仅迭代单个样本也会抛出相同错误。

相关代码片段 #

数据集特征定义 #

features=tfds.features.FeaturesDict({
    # 数据集特征定义：图像、位置、动作、提示
    'image': tfds.features.Image(shape=(320, 180, 1)),
    'position': tfds.features.Tensor(shape=(2,), dtype=np.float32),
    'action': tfds.features.Scalar(dtype=np.int8),
    'hint': tfds.features.Scalar(dtype=np.bool_),
}),
supervised_keys=(('image', 'hint'), ('position', 'action')),

数据生成逻辑 #

for row in csv.DictReader(csvData):
    fileName=row['fileName']
    yield fileName, {
        'image': path+fileName,
        'position' : [np.float32(row['Xpos']), np.float32(row['Ypos'])],
        'action' : np.int8(row['action']),
        'hint' : tfds.core.utils.bool_utils.parse_bool(row['hint']),
    }

数据集加载代码 #

dataset_train=tfds.load('my_dataset', split='all', batch_size=64, as_supervised=True)
for example in dataset_train.take(1):    pass

问题原因及解决办法 #

核心原因 #

tfds.features.Image默认会将图像以JPEG/PNG格式编码存储，加载解码时可能自动转换通道数：比如单通道JPEG图像没有明确的通道标识，解码后会被默认转为3通道RGB，导致实际shape为(320,180,3)，与定义的(320,180,1)冲突。TensorFlow在批量处理时尝试对shape不匹配的样本做padding，就会触发"填充到更小尺寸"的错误。

具体解决步骤 #

1. 修正Image特征参数
   显式指定图像的编码格式和数据类型，确保解码后shape与定义一致：

   'image': tfds.features.Image(shape=(320, 180, 1), dtype=tf.uint8, encoding_format='png'),
   

   注意：仅当图像为PNG格式时，encoding_format='png'才能确保单通道解码；若使用JPEG格式，建议先将图像转换为单通道PNG，或改用手动读取方式。

2. 手动读取并处理图像
   放弃tfds.features.Image的自动加载，手动读取图像并强制转成单通道：

   • 修改数据生成逻辑，直接返回图像张量：

     for row in csv.DictReader(csvData):
         fileName=row['fileName']
         image_path = path + fileName
         # 手动读取并解码图像
         image = tf.io.read_file(image_path)
         image = tf.image.decode_png(image, channels=1)  # 指定单通道解码
         image = tf.cast(tf.image.resize(image, (320, 180)), tf.uint8)  # 确保尺寸一致
         yield fileName, {
             'image': image,
             'position' : [np.float32(row['Xpos']), np.float32(row['Ypos'])],
             'action' : np.int8(row['action']),
             'hint' : tfds.core.utils.bool_utils.parse_bool(row['hint']),
         }
     

   • 同时将features中的image定义改为Tensor类型：

     'image': tfds.features.Tensor(shape=(320, 180, 1), dtype=tf.uint8),
     

3. 排查样本shape一致性
   先关闭批量加载，逐个检查样本的shape是否符合定义：

   dataset_train=tfds.load('my_dataset', split='all', batch_size=None, as_supervised=True)
   for x, y in dataset_train.take(5):
       print("图像shape:", x[0].shape)
       print("提示shape:", x[1].shape)
       print("位置shape:", y[0].shape)
       print("动作shape:", y[1].shape)
   

   若发现个别样本shape异常，需排查对应图像文件是否存在尺寸或通道数问题。

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