如何像PyTorch那样通过子类化tf.data.Dataset创建TensorFlow自定义数据集？

TensorFlow官方并不推荐子类化tf.data.Dataset来实现自定义数据集——它的底层基于计算图设计，子类化需要处理大量图兼容细节，反而会让代码复杂难调试。以下两种方案更符合TensorFlow设计理念，能满足你模块化、灵活处理数据加载/预处理/增强的需求：

方法一：组合式API（官方推荐） #

通过拆分数据加载、预处理、增强为独立函数，再用tf.data.Dataset的链式API组合，结构清晰且天然支持图模式、分布式训练与性能优化。

示例代码（对应人脸关键点数据集场景） #

import tensorflow as tf
import pandas as pd
import numpy as np

# 1. 定义数据加载逻辑
def load_sample(csv_row):
    # 拼接图片路径并读取
    image_path = tf.strings.join([csv_row['root_dir'], csv_row['image_path']])
    image = tf.io.read_file(image_path)
    image = tf.image.decode_jpeg(image, channels=3)
    # 解析关键点数据
    landmarks = tf.cast(csv_row['landmarks'], tf.float32)
    landmarks = tf.reshape(landmarks, (-1, 2))
    return image, landmarks

# 2. 定义预处理逻辑
def preprocess(image, landmarks):
    # 统一图片尺寸
    image = tf.image.resize(image, (224, 224))
    # 图片归一化
    image = tf.cast(image, tf.float32) / 255.0
    return image, landmarks

# 3. 定义数据增强逻辑
def augment(image, landmarks):
    # 随机水平翻转（同步处理关键点）
    if tf.random.uniform(()) > 0.5:
        image = tf.image.flip_left_right(image)
        landmarks = tf.stack([1.0 - landmarks[:, 0], landmarks[:, 1]], axis=-1)
    return image, landmarks

# 4. 构建完整数据管道
def create_face_landmarks_dataset(csv_file, root_dir, batch_size=32, shuffle=True):
    # 读取CSV并转换为TensorFlow数据集
    df = pd.read_csv(csv_file)
    dataset = tf.data.Dataset.from_tensor_slices({
        'image_path': df['image_path'].values,
        'landmarks': df['landmarks'].apply(lambda x: np.array(x.split(','), dtype=np.float32)).values,
        'root_dir': [root_dir] * len(df)
    })

    # 链式组合所有操作，启用多线程加速
    dataset = dataset.map(load_sample, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.map(augment, num_parallel_calls=tf.data.AUTOTUNE)

    # 洗牌、分批、预取优化
    if shuffle:
        dataset = dataset.shuffle(buffer_size=len(df))
    dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return dataset

# 使用示例
train_dataset = create_face_landmarks_dataset('train.csv', './face_images', batch_size=32)
for images, landmarks in train_dataset.take(1):
    print(f"Batch images shape: {images.shape}, Batch landmarks shape: {landmarks.shape}")

方法二：自定义生成器类（贴近PyTorch写法） #

如果你更习惯类封装的方式，可以用生成器类封装所有逻辑，再通过tf.data.Dataset.from_generator转换为TensorFlow数据集。这种写法和PyTorch的Dataset子类模式高度相似，适合快速原型开发。

示例代码 #

import tensorflow as tf
import pandas as pd
import os
from PIL import Image
import numpy as np

class FaceLandmarksDataset:
    def __init__(self, csv_file, root_dir, transform=None):
        self.df = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # 加载单样本数据
        image_path = os.path.join(self.root_dir, self.df.iloc[idx]['image_path'])
        image = Image.open(image_path).convert('RGB')
        landmarks = np.array(self.df.iloc[idx]['landmarks'].split(','), dtype=np.float32).reshape(-1, 2)

        # 应用自定义变换
        if self.transform:
            image, landmarks = self.transform(image, landmarks)

        # 转换为TensorFlow张量
        image = tf.convert_to_tensor(np.array(image), dtype=tf.float32) / 255.0
        landmarks = tf.convert_to_tensor(landmarks, dtype=tf.float32)
        return image, landmarks

# 将自定义类转换为tf.data.Dataset
def create_dataset_from_class(csv_file, root_dir, batch_size=32, shuffle=True):
    dataset_instance = FaceLandmarksDataset(csv_file, root_dir)
    # 指定输出张量的类型和形状（必须明确，否则TensorFlow无法优化）
    output_signature = (
        tf.TensorSpec(shape=(None, None, 3), dtype=tf.float32),
        tf.TensorSpec(shape=(68, 2), dtype=tf.float32)
    )
    tf_dataset = tf.data.Dataset.from_generator(
        lambda: (dataset_instance[i] for i in range(len(dataset_instance))),
        output_signature=output_signature
    )

    # 数据管道优化
    if shuffle:
        tf_dataset = tf_dataset.shuffle(buffer_size=len(dataset_instance))
    tf_dataset = tf_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return tf_dataset

# 使用示例
train_dataset = create_dataset_from_class('train.csv', './face_images', batch_size=32)
for images, landmarks in train_dataset.take(1):
    print(f"Batch images shape: {images.shape}, Batch landmarks shape: {landmarks.shape}")

方案对比 #

• 组合式API：完全基于TensorFlow图操作，性能最优，支持分布式训练和自动微分，适合生产环境。
• 生成器类：写法贴近PyTorch，学习成本低，但依赖Python生成器，在图模式下可能存在性能损耗，适合快速验证想法。

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