使用 Hugging Face Transformers 微调视觉 Transformer (ViT) 进行图像分类

概览

Hugging Face 已详细说明了用于图像分类任务的 Vision Transformer (ViT) 微调工作流程。通过将图像补丁视为标记——类似于自然语言处理 (NLP) 中如何处理单词——ViT 使基于 transformer 的架构能够应用于计算机视觉。此过程涉及将图像分割成一个子图像块的网格，使用线性投影嵌入每个补丁，并将得到的标记序列传递给 transformer 模型。

使用 ViTImageProcessor 进行图像预处理

正确的图像变换对 ViT 模型性能至关重要，因为模型期望输入遵循其原始训练期间使用的特定归一化和调整大小参数。

为了保持一致性，ViTImageProcessor 用于从预训练模型（如 google/vit-base-patch16-224-in21k）加载配置。处理器执行以下操作：

• 调整大小：将图像调整为标准尺寸（例如，224x224 像素）。
• 归一化：对像素数据应用特定的均值和标准差值。
• 重采样：确保图像正确采样以供模型使用。

通过此工具处理图像将返回一个包含 pixel_values 的字典，这是模型所需的数值张量表示。

使用 Transforms 进行高效数据集处理

在处理图像数据集时，通过 ds.map 对每个示例应用变换可能在计算上昂贵且缓慢。相反，Hugging Face 建议使用 datasets 库的 with_transform 方法。

变换在索引示例时实时应用，而不是预处理整个数据集。这种方法需要一个能够处理批量数据的变换函数，将 PIL 图像列表转换为所需的 pixel_values 张量，同时保留相关标签。

微调流程和配置

微调 ViT 模型涉及几个关键组件，以确保模型收敛并正确评估：

数据整理和指标

由于批次以字典列表的形式交付，需要自定义 collate_fn 将 pixel_values 和 labels 堆叠为 torch 张量。对于评估，通常使用 evaluate 库中的 accuracy 指标来比较模型的预测类别（通过在预测上使用 np.argmax 确定）与真实标签。

模型初始化

为了将预训练的 ViT 调整为特定任务，使用 ViTForImageClassification。模型通过 num_labels 初始化以创建具有正确输出单元数量的分类头。此外，提供 id2label 和 label2id 映射以确保模型在 Hugging Face Hub 上托管时的输出可读。

训练参数

TrainingArguments 中的关键配置包括：

• remove_unused_columns=False：这一点至关重要。默认情况下，训练器会删除模型前向传播未使用的列。然而，由于 image 列被变换函数用于创建 pixel_values，因此必须保留。
• fp16=True：启用混合精度训练以降低内存使用并加速训练。
• evaluation_strategy="steps"：允许在训练过程中定期进行评估。

性能结果

在使用 beans 数据集（用于区分健康与不健康的豆叶）的提供示例中，微调后的 ViT 模型在 4 个训练周期后达到了评估准确率 0.985 和评估损失 0.0637。