如何在自有数据集上使用TensorFlow CycleGAN？自定义数据加载方法问询

完全可以实现类似tfds.load的效果，用TensorFlow的tf.data API就能从你的本地文件夹生成格式一致的Dataset对象，不需要依赖flow_from_directory（当然也可以用，但tf.data更贴合官方教程的风格）。下面是具体的实现步骤：

1. 定义通用的图片加载与预处理函数 #

首先要对齐官方教程中对horse2zebra数据集的预处理逻辑，保证输入数据的格式、归一化方式和数据增强策略完全一致，这样模型训练才能和教程预期的效果对齐：

import tensorflow as tf

def load_and_preprocess_image(image_path, img_size=256, crop_size=256):
    # 加载图片
    img = tf.io.read_file(image_path)
    img = tf.image.decode_jpeg(img, channels=3)
    # 转换为float32并归一化到[-1, 1]（CycleGAN的输入要求）
    img = tf.cast(img, tf.float32)
    img = (img / 127.5) - 1.0
    # 数据增强：随机裁剪+随机水平翻转（和教程一致）
    img = tf.image.random_flip_left_right(img)
    img = tf.image.resize(img, [img_size, img_size])
    img = tf.image.random_crop(img, [crop_size, crop_size, 3])
    # 返回(image, label)元组（和tfds.load的as_supervised=True格式一致，标签用0占位即可）
    return img, 0

def load_dataset_from_dir(dir_path, batch_size=1):
    # 获取文件夹下所有图片路径
    dataset = tf.data.Dataset.list_files(f"{dir_path}/*")
    # 映射预处理函数
    dataset = dataset.map(load_and_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)
    # 打乱、批处理、预取
    dataset = dataset.shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return dataset

2. 加载你的自定义数据集 #

现在就可以像你期望的那样，一行代码加载单个文件夹的数据集，和教程里的train_horses格式完全一致：

# 加载训练集
train_horses = load_dataset_from_dir("data/trainA")
train_zebras = load_dataset_from_dir("data/trainB")

# 加载测试集（测试集不需要数据增强，所以单独写个简化版函数）
def load_test_dataset_from_dir(dir_path, batch_size=1):
    dataset = tf.data.Dataset.list_files(f"{dir_path}/*")
    def test_preprocess(image_path):
        img = tf.io.read_file(image_path)
        img = tf.image.decode_jpeg(img, channels=3)
        img = tf.cast(img, tf.float32)
        img = (img / 127.5) - 1.0
        img = tf.image.resize(img, [256, 256])
        return img, 0
    dataset = dataset.map(test_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return dataset

test_horses = load_test_dataset_from_dir("data/testA")
test_zebras = load_test_dataset_from_dir("data/testB")

3. 验证格式一致性 #

你可以打印一下数据集的元素，确认和tfds.load返回的格式完全相同：

for img, label in train_horses.take(1):
    print(f"Image shape: {img.shape}, Label: {label}")

输出应该类似Image shape: (1, 256, 256, 3), Label: 0，和官方教程里的train_horses结构一致，后续的模型训练代码可以直接复用。

补充说明 #

• 如果你的图片格式不是JPEG，可以把tf.image.decode_jpeg改成tf.image.decode_png或者tf.image.decode_image（自动识别格式）。
• 如果你需要和Keras的flow_from_directory类似的功能（比如自动分类标签），也可以基于tf.data实现，但CycleGAN本身不需要标签，所以上面的 dummy 标签完全够用。

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