使用zip对象传入model.fit()报错：数据类型无法识别的解决方法

问题背景 #

开发分割模型时，使用ImageDataGenerator()加载图像和掩码并返回zip对象，传入model.fit()时触发如下错误：

ValueError: Unrecognized data type: x=<zip object at 0x00000172364F6040> (of type <class 'zip'>)

原代码如下：

def create_segmentation_generator_train(img_path, msk_path, BATCH_SIZE):
    data_gen_args = #augmentation code
    datagen = ImageDataGenerator(**data_gen_args)
    
    img_generator = datagen.flow_from_directory(img_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    msk_generator = datagen.flow_from_directory(msk_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    return zip(img_generator, msk_generator)

# Remember not to perform any image augmentation in the test generator!
def create_segmentation_generator_test(img_path, msk_path, BATCH_SIZE):
    data_gen_args = dict(rescale=1./255)
    datagen = ImageDataGenerator(**data_gen_args)
    
    img_generator = datagen.flow_from_directory(img_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    msk_generator = datagen.flow_from_directory(msk_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    return zip(img_generator, msk_generator)

train_generator = create_segmentation_generator_train(data_dir_train_image, data_dir_train_mask, BATCH_SIZE_TRAIN)
test_generator = create_segmentation_generator_test(data_dir_test_image, data_dir_test_mask, BATCH_SIZE_TEST)

解决方案 #

Keras不支持直接将原生zip对象作为训练输入，需将其包装为Keras可识别的生成器或tf.data.Dataset，以下是两种可行方案：

方案1：包装为自定义生成器 #

修改生成器函数，返回一个逐批次输出图像和掩码对的Python生成器：

def create_segmentation_generator_train(img_path, msk_path, BATCH_SIZE):
    data_gen_args = # 填入你的图像增强参数
    datagen = ImageDataGenerator(**data_gen_args)
    
    img_generator = datagen.flow_from_directory(img_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    msk_generator = datagen.flow_from_directory(msk_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    
    def seg_generator():
        for img_batch, msk_batch in zip(img_generator, msk_generator):
            # 掩码归一化到0-1区间（多分类任务需转one-hot编码）
            msk_batch = msk_batch / 255.0
            yield img_batch, msk_batch
    return seg_generator()

# 测试生成器同步修改
def create_segmentation_generator_test(img_path, msk_path, BATCH_SIZE):
    data_gen_args = dict(rescale=1./255)
    datagen = ImageDataGenerator(**data_gen_args)
    
    img_generator = datagen.flow_from_directory(img_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    msk_generator = datagen.flow_from_directory(msk_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    
    def seg_generator():
        for img_batch, msk_batch in zip(img_generator, msk_generator):
            msk_batch = msk_batch / 255.0
            yield img_batch, msk_batch
    return seg_generator()

方案2：使用tf.data.Dataset包装 #

借助TensorFlow的Dataset API整合生成器，适配性更强且性能更优：

import tensorflow as tf

def create_segmentation_generator_train(img_path, msk_path, BATCH_SIZE):
    data_gen_args = # 填入你的图像增强参数
    datagen = ImageDataGenerator(**data_gen_args)
    
    img_generator = datagen.flow_from_directory(img_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    msk_generator = datagen.flow_from_directory(msk_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    
    # 转换为tf.data.Dataset并指定输出格式
    seg_dataset = tf.data.Dataset.from_generator(
        lambda: zip(img_generator, msk_generator),
        output_signature=(
            tf.TensorSpec(shape=(BATCH_SIZE, *IMG_SIZE, 1), dtype=tf.float32),
            tf.TensorSpec(shape=(BATCH_SIZE, *IMG_SIZE, 1), dtype=tf.float32)
        )
    )
    # 添加掩码预处理
    seg_dataset = seg_dataset.map(lambda x, y: (x, y / 255.0))
    return seg_dataset

# 测试生成器同步修改
def create_segmentation_generator_test(img_path, msk_path, BATCH_SIZE):
    data_gen_args = dict(rescale=1./255)
    datagen = ImageDataGenerator(**data_gen_args)
    
    img_generator = datagen.flow_from_directory(img_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    msk_generator = datagen.flow_from_directory(msk_path, target_size=IMG_SIZE, class_mode=None, color_mode='grayscale', batch_size=BATCH_SIZE, seed=SEED)
    
    seg_dataset = tf.data.Dataset.from_generator(
        lambda: zip(img_generator, msk_generator),
        output_signature=(
            tf.TensorSpec(shape=(BATCH_SIZE, *IMG_SIZE, 1), dtype=tf.float32),
            tf.TensorSpec(shape=(BATCH_SIZE, *IMG_SIZE, 1), dtype=tf.float32)
        )
    )
    seg_dataset = seg_dataset.map(lambda x, y: (x, y / 255.0))
    return seg_dataset

注意事项 #

• 掩码预处理需根据任务调整：二分类分割归一化到0-1即可，多分类需转换为one-hot编码。
• output_signature中的形状要与你的IMG_SIZE、批量大小和通道数匹配。

内容的提问来源于Stack Exchange，提问作者Ashwin Singh