在Keras/TensorFlow中使用Hugging Face ImageNet-1k数据集报错求助

问题根源 #

你遇到的ValueError核心原因是：Hugging Face的Dataset对象即便设置了tensorflow格式，本质仍是Arrow Dataset类型，而Keras的fit()方法仅识别TensorFlow原生的tf.data.Dataset、numpy数组、列表等数据类型，无法直接处理Hugging Face的Dataset实例。

解决方案：转换为TensorFlow原生Dataset #

使用Hugging Face Dataset提供的to_tf_dataset()方法，将数据集转换为Keras可识别的格式，同时完成自编码器所需的输入/标签配对（输入和标签为同一图像）。

修改后的完整代码 #

from datasets import load_dataset
import tensorflow as tf

# 加载ImageNet-1k数据集
ds = load_dataset('imagenet-1k')

# 定义预处理函数：适配模型输入，生成自编码器的输入/标签对
def preprocess_sample(sample):
    # 将PIL图像转为TF张量，调整到模型要求的尺寸（示例为224x224，按需修改）
    img_tensor = tf.convert_to_tensor(sample['image'])
    img_resized = tf.image.resize(img_tensor, (224, 224))
    # 归一化到[0,1]区间（可选，根据模型需求调整）
    img_normalized = tf.cast(img_resized, tf.float32) / 255.0
    # 自编码器的输入和标签为同一图像
    return img_normalized, img_normalized

# 转换训练集为TF Dataset
train_ds = ds['train'].to_tf_dataset(
    columns=['image'],          # 指定输入列
    label_cols=['image'],       # 指定标签列（与输入列一致）
    shuffle=True,               # 训练集开启打乱
    batch_size=32,              # 根据GPU显存调整批量大小
    collate_fn=preprocess_sample # 应用预处理函数
)

# 转换验证集为TF Dataset
val_ds = ds['validation'].to_tf_dataset(
    columns=['image'],
    label_cols=['image'],
    shuffle=False,              # 验证集无需打乱
    batch_size=32,
    collate_fn=preprocess_sample
)

# 训练自编码器
autoencoder.fit(
    train_ds,
    epochs=10,
    validation_data=val_ds
)

替代方案：先预处理再转换 #

若更习惯用Dataset的map()方法做预处理，可先处理再转换为TF Dataset：

# 先对数据集做预处理映射
processed_train = ds['train'].map(
    lambda x: {
        'image': tf.cast(tf.image.resize(tf.convert_to_tensor(x['image']), (224,224)), tf.float32) / 255.0
    }
)

# 转换为TF Dataset，自动生成输入-标签对
train_ds = processed_train.to_tf_dataset(
    columns=['image'],
    label_cols=['image'],
    shuffle=True,
    batch_size=32
)

关键注意事项 #

• 预处理时需确保输出张量的shape与自编码器输入层shape完全匹配
• 批量大小batch_size需根据GPU显存调整，避免内存溢出
• 验证集无需开启shuffle，训练集开启shuffle可提升模型泛化性

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