image_dataset_from_directory适配TFLite量化representative_dataset

原有代码失效的核心原因 #

• image_dataset_from_directory生成的数据集迭代时，每个元素是**(图像批次张量, 标签批次张量)**的二元组。原实现直接对整个遍历结果做类型转换、直接输出，既把冗余的标签数据传入量化流程，也不符合representative_dataset仅输出模型输入张量的格式要求。
• 原train_ds本身已经按照配置的batch_size完成了分批次操作，直接在其上调用.batch(1)会在原有批次维度外层嵌套新的批量维度，最终输入张量形状为(1, 原batch_size, 图像高, 图像宽, 通道数)，和模型要求的单样本(1, 图像高, 图像宽, 通道数)输入形状不匹配，触发维度校验报错。

可正常运行的适配代码 #

实现时需要先拆分原数据集的批次维度，恢复单样本粒度的数据流，再按照量化要求提取图像部分、做类型转换输出：

def representative_dataset():
    # 拆分原有批次为单样本粒度，再按1样本批量分组，取前100个样本
    for image, _ in train_ds.unbatch().batch(1).take(100):
        # 仅提取图像张量转float32格式，按要求封装为列表输出
        yield [image.astype(tf.float32)]

实现注意点 #

• 该逻辑会自动继承原数据集上配置的所有图像预处理、归一化、性能优化（缓存、预取）操作，无需额外调整，可保证量化时的输入分布和模型实际推理的输入分布完全一致。
• 无需修改原train_ds的构建参数，unbatch操作仅临时拆分批次用于生成代表数据集，不会改动原训练数据集的结构。

内容的提问来源于stack exchange，提问作者neil.morrison