使用ViT+Keras/Transformers做图像分类时遇KeyError: 'img'求助

错误原因 #

用load_dataset("imagefolder")加载数据集后，图像字段的默认键名是**'image'**，但你写的augmentation和process函数里错误用了'img'作为键名，导致找不到对应字段抛出KeyError。

修正方案 #

把两个处理函数里的键名从'img'改成'image'，同时注意Keras数据增强模块需要输入张量，要把PIL格式的图片转成张量：

修正后的图像增强函数

def augmentation(examples):
    examples["pixel_values"] = [data_augmentation(tf.convert_to_tensor(image)) for image in examples["image"]]
    return examples

修正后的基础处理函数

def process(examples):
    examples.update(feature_extractor(examples['image']))
    return examples

修正后的关键代码片段 #

# 图像增强处理
def augmentation(examples):
    examples["pixel_values"] = [data_augmentation(tf.convert_to_tensor(image)) for image in examples["image"]]
    return examples

# 基础预处理
def process(examples):
    examples.update(feature_extractor(examples['image']))
    return examples

# 重命名标签字段适配后续流程
dataset_ds = dataset["train"].rename_column("label", "labels")

# 执行数据处理
processed_dataset = dataset_ds.map(augmentation, batched=True)

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