Hugging Face 图像分类：示例详解全指南

在快速发展的人工智能（AI）与机器学习（ML）领域，图像分类是一项基础性任务，支撑着从人脸识别到医学影像诊断等广泛应用。

在众多面向数据科学家与 AI 从业者的工具与库中，Hugging Face 一直扮演着重要角色，尤其是在自然语言处理（NLP）以及日益增长的计算机视觉任务（如图像分类）方面。

本文将讲解使用 Hugging Face 进行图像分类的要点，包括理解图像分类基础、准备数据、训练模型、将模型部署到 Hugging Face Hub，以及通过 API 和 Hugging Face 界面与已部署模型交互。

无论您是初学者还是有经验的实践者，本指南都将帮助您高效利用 Hugging Face 满足图像分类需求。

借助 Hugging Face 理解图像分类

本节将介绍图像分类基础以及使用 Hugging Face 的优势。

图像分类基础

图像分类是指将图像归入若干预定义类别之一。

该任务通过算法分析图像的视觉内容，并基于从训练数据中学到的模式预测其类别。由于能够捕捉图像数据中的相关模式，深度学习模型，尤其是卷积神经网络（CNN），已成为此任务的标准方法。

若想进一步了解 CNN，可参考我们的文章 An Introduction to Convolutional Neural Networks (CNNs)，系统讲解 CNN 的原理、其对图像分析的影响，以及在深度学习应用中提升鲁棒性、对抗过拟合的一些关键策略。

此外，Classification in Machine Learning: An Introduction 一文以分步方式介绍机器学习中的分类：它是什么、如何使用，以及一些分类算法示例。

为何选择 Hugging Face 进行图像分类？

Hugging Face 提供了完善的机器学习生态系统，界面易用、拥有海量预训练模型库，并提供完善的文档。

与其他平台相比，Hugging Face 在图像分类方面具有以下优势：

使用 Hugging Face 的三大优势（图像由 GPT-4 生成）

易用性：Hugging Face 的用户友好型 API 与完整文档，让初学者与专家都能轻松上手。
预训练模型：可访问范围广泛的预训练模型，并在自定义数据集上进行微调，从而节省时间与算力资源。任何注册用户都可以训练并部署自己的模型。
社区与支持：活跃的社区与支持论坛，便于排查问题、持续改进模型。

此外，Hugging Face 支持将模型部署到 AWS、Azure、Google Cloud Platform 等主流云平台，并提供多种推理类型。

云平台上的模型部署选项

为图像分类准备数据

本文的核心任务是图像分类，我们将使用 Hugging Face 的 beans 数据集作为示例。加载数据集后，我们会先进行可视化，再进入训练前的预处理。

配套代码笔记本可在 Google Colab 中跟练。

本文中的代码主要参考自 Hugging Face 官网。

依赖库要求

本文使用了多种库，各司其职。其中部分库需要安装。可通过 Python 包管理器“pip”按如下方式安装：

%%bash
pip -q install datasets
pip -q install transformers=='4.29.0'
pip -q install tensorflow=='2.15' # At least this tensorflow version is required to use the "evaluate module"
pip -q install evaluate
pip -q install --upgrade accelerate技术支持

安装完成后，务必重启内核。成功安装后，即可通过 from 与 import 语句导入所需库：

import torch
import torchvision
import numpy as np
import evaluate
from datasets import load_dataset
from huggingface_hub import notebook_login
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from transformers import DefaultDataCollator
from transformers import AutoImageProcessor
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
from transformers import AutoModelForImageClassification, TrainingArguments, Trainer
import matplotlib.pyplot as plt技术支持

数据需求与组织

使用 load_dataset 函数即可便捷地从 Hugging Face 加载数据。

beans_train = load_dataset("beans", split="train")技术支持

beans 是所使用的数据集名称。
train 表示所用的是整个数据集中的训练子集。

数据加载后，只需打印上述变量名即可查看其特性。

print(beans_train)技术支持

输出如下：

Dataset({
    features: ['image_file_path', 'image', 'labels'],
    num_rows: 1034
})技术支持

共有 1034 张图像，每张图像包含三个字段：

image_path：图像的实际路径。
image：豆叶图像的 PIL 对象。
labels：图像标签，可能为以下三类之一：

0: angular_leaf_spot
1: bean_rust
2: healthy

下面的辅助函数便于随机可视化若干图像，并展示其对应的标签类型与标签编号（0、1、2）。

labels_names = {
   0: "angular_leaf_spot",
   1: "bean_rust",
   2: "healthy",
}

def display_random_images(dataset, num_images=4):

   num_rows = 2
   num_cols = np.ceil(num_images / num_rows).astype(int)

   plt.figure(figsize=(num_cols * 3, num_rows * 3))
   indices = np.random.choice(range(len(dataset)), size=num_images, replace=False)

   for i, idx in enumerate(indices, 1):
       idx = int(idx)
       image = dataset[idx]['image']
       label = dataset[idx]['labels']
       label_name = labels_names.get(label, "Unknown")

       plt.subplot(num_rows, num_cols, i)
       plt.imshow(image)
       plt.title(f"{label_name} ({label})")
       plt.axis('off')

   plt.tight_layout()
   plt.show()技术支持

简而言之，display_random_images 会从给定数据集中随机选取一组图像并以网格形式展示，每张图像标注其分类（如 "angular_leaf_spot"、"bean_rust"、"healthy"）。

函数会将图像动态排布为两行，并根据要展示的图像总数计算所需列数；确保不重复选图，并为每张图像显示对应的标签名与编号。

随后我们可按如下方式显示 6 张随机图像：

display_random_images(beans_train, num_images=6)技术支持

数据集中随机六张图像

数据预处理

在训练模型前，需要进行适当的预处理，以便数据符合模型要求。

首先，将原始数据集按 80%/20% 切分为训练集与验证集，使用 train_test_split 实现。

beans_train = beans_train.train_test_split(test_size=0.2)技术支持

接着，创建两个字典：label2id 与 id2label。这对于在人类可读标签名与机器可读数值 ID 之间进行转换至关重要。

该过程既有助于训练，也便于解读最终模型的预测结果。

labels = beans_train["train"].features["labels"].names
label2id, id2label = dict(), dict()

for i, label in enumerate(labels):
   label2id[label] = str(i)
   id2label[str(i)] = label技术支持

如下是 id2label 变量的结果：

print(id2label)技术支持

{'0': 'angular_leaf_spot', '1': 'bean_rust', '2': 'healthy'}技术支持

加载预训练模型

我们的分类模型基于 Hugging Face 的预训练模型 Vision Transformer（ViT）。

该模型发表于论文 An Image is Worth 16x16 words: Transformers for Image Recognition at Scale，这是 Transformer 架构首次成功应用于图像识别，并在 ImageNet 数据集上完成训练。

以下代码展示了加载预训练模型、对数据进行符合模型要求的预处理、以及确保图像正确处理并可输入预训练模型以进行微调与评估的各步骤。

checkpoint = "google/vit-base-patch16-224-in21k"
image_processor = AutoImageProcessor.from_pretrained(checkpoint)

normalize = Normalize(mean=image_processor.image_mean,
                  	std=image_processor.image_std)

size = (
	image_processor.size["shortest_edge"]
	if "shortest_edge" in image_processor.size
	else (image_processor.size["height"], image_processor.size["width"])
)

_transforms = Compose([RandomResizedCrop(size), ToTensor(), normalize])

def transforms(examples):
 
	examples["pixel_values"] = [_transforms(img.convert("RGB")) for img in examples["image"]]
	del examples["image"]
	return examples

beans_transformed = beans_train.with_transform(transforms)

data_collator = DefaultDataCollator()技术支持

我们逐行解读上述代码：

首先，从 google/vit-base-patch16-224-in21k 加载预训练图像处理模型 checkpoint。
随后为该模型创建图像处理器，包含归一化参数与期望的图像尺寸。
基于图像处理器给出的均值与标准差，定义归一化变换。
确定用于变换的图像尺寸，优先取最短边，否则使用给定的高宽。
组合一系列变换（_transforms），包括随机裁剪、张量转换与归一化。
定义 transforms 辅助函数，将这些变换应用于输入样本的 "image" 字段，将图像转换为 RGB，并移除原始的 "image" 字段。
对 beans_train 数据集应用 transforms 函数，为模型训练做好准备。
最后，实例化默认的数据整理器，以便在训练中进行批处理。

完成以上设置后，我们使用准确率作为评估指标，并通过以下辅助函数实现：

accuracy = evaluate.load("accuracy")
def compute_metrics(eval_pred):
   predictions, labels = eval_pred
   predictions = np.argmax(predictions, axis=1)

   return accuracy.compute(predictions=predictions,
                           references=labels)技术支持

使用 Hugging Face 训练模型

在开始训练前，利用 .from_pretrained 基于预训练模型定义实际要训练的模型：

model = AutoModelForImageClassification.from_pretrained(
   checkpoint,
   num_labels=len(labels),
   id2label=id2label,
   label2id=label2id,
)技术支持

微调与评估

下一步是通过配置并初始化训练环境来设置模型。

该设置包括指定与模型训练、评估与优化相关的关键超参数。

在此之前，需要使用账户凭据登录 Hugging Face 账号，获取方式如下：

获取 Hugging Face 凭据的主要步骤

notebook_login()技术支持

根据提示，粘贴上面步骤 4 中复制的连接令牌。

登录成功后，应显示如下信息：

以下是配置与初始化代码：

training_args = TrainingArguments(
	output_dir="./beans_health_type_classifier",
	remove_unused_columns=False,
	evaluation_strategy="epoch",
	save_strategy="epoch",
	learning_rate=5e-5,
	per_device_train_batch_size=16,
	gradient_accumulation_steps=4,
	per_device_eval_batch_size=16,
	num_train_epochs=3,
	warmup_ratio=0.1,
	logging_steps=10,
	load_best_model_at_end=True,
	metric_for_best_model="accuracy"
)

trainer = Trainer(
	model=model,
	args=training_args,
	data_collator=data_collator,
	train_dataset=beans_transformed["train"],
	eval_dataset=beans_transformed["test"],
	tokenizer=image_processor,
	compute_metrics=compute_metrics,
)技术支持

通过 TrainingArguments 配置训练参数，如输出目录、按 epoch 进行评估与保存、学习率、批大小、梯度累积步数等，以优化训练过程与模型表现。
用上述模型与参数初始化 Trainer，并指定数据整理器、训练与评估数据集、用于预处理的数据处理器（此处为图像处理器），以及计算评估指标的函数。

接下来进入训练阶段，使用 .train 方法启动。

trainer.train()技术支持

本次训练得到的指标如下：

三轮训练的模型结果

可以看到，表格展示了 beans 数据集上三轮训练的关键指标：训练损失、验证损失与准确率。

训练损失：从第一轮的 1.02800 降至第三轮的 0.53470，表明模型在持续学习，训练集表现不断提升。
验证损失：从第一轮的 0.77973 降至第三轮的 0.43553，说明模型泛化良好，未对训练数据过拟合。
准确率：从第一轮的 85.5072% 稳步提升至第三轮的 94.2029%，反映出模型在验证集上对豆叶图像的分类能力明显增强。

部署您的图像分类模型

首次在 Hugging Face 门户进行认证时，连接模式为只读。但要将模型推送到 Hugging Face Hub，需要具有写入权限的凭据。

之后，使用新创建的凭据再次通过 notebook_login 登录。

notebook_login()技术支持

使用 push_to_hub 方法将模型推送到 Hub：

trainer.model.push_to_hub("zoumana/beans_health_type_classifier")技术支持

模型将被推送至 Hub，您可以查看其主要属性。

模型已推送至 Hugging Face

集成到应用

模型共享后，可通过 REST API 集成到应用中，或直接在 Hugging Face 门户上调用。开发者与研究人员也可以通过 Transformers 库加载并使用该模型。

本节介绍三种主要使用方式。

1. 在 Hugging Face 门户中使用

在门户中，用户可以直接上传图像并获取模型预测，如下动图所示。

在 Hugging Face 门户中使用模型

2. 在 Transformers 中使用

点击“Use in Transformers”图标，即可自动获取如下使用代码：

结合 Transformers 使用模型

3. 使用 REST API

最后，可通过 REST API 使用模型。在门户中点击“Deploy > Inference API (serverless)”图标，即可自动生成连接字符串。

使用 API 端点

上方动图中的代码片段如下所示，可用于对指定图像进行预测。

所用图像与在 Hugging Face 门户演示时相同。

请将“xxx”替换为实际值。

import requests

API_URL = "https://api-inference.huggingface.co/models/zoumana/beans_health_type_classifier"
headers = {"Authorization": "Bearer xxxxxxxxxxxxxxxxx"}

def query(filename):
   with open(filename, "rb") as f:
       data = f.read()
   response = requests.post(API_URL, headers=headers, data=data)
   return response.json()技术支持

随后，使用上述 query 函数获取模型响应：

output = query("./inference_images/test_beans_leaf.jpg")技术支持

最终结果为一个 JSON 响应，如下所示。可见概率最高的标签即为预测结果；在此示例中为 healthy，与 Hugging Face 门户上的预测一致。

print(output)技术支持

[
    {'label': 'healthy', 'score': 0.4799719452857971},
     {'label': 'angular_leaf_spot', 'score': 0.27489492297172546},
     {'label': 'bean_rust', 'score': 0.2451331913471222}
]技术支持

结语与延伸学习

本文完整展示了如何使用 Hugging Face 开展图像分类任务。

我们从图像分类基础出发，说明了为什么 Hugging Face 借助其丰富的模型库与良好的社区支持，是完成此任务的理想选择。

随后介绍了如何为分类准备训练图像，包括如何确保训练数据满足微调需求。

接着，我们讲解了如何选择合适的预训练模型并进行微调，并介绍了将训练后模型集成进真实应用的三种主要方式。

AI 领域始终在快速演进，您的每一步实践都会打开新的可能。若您希望进一步学习 AI 与机器学习，特别是图像分类及相关主题，推荐参考以下资源：

面向初学者的 Transformers 与 Hugging Face 入门，请参阅 An Introduction to Using Transformers and Hugging Face。
若想在 Python 中更深入地处理图像，可学习 Image Processing with Python 技能路径，了解如何处理与分析图像。
若需概览图像识别的概念与重要性，请阅读 What is Image Recognition?。

本指南旨在帮助初级、中级到高级实践者使用 Hugging Face 开展图像分类项目。

作者

Zoumana Keita

LinkedIn Twitter

一位多才多艺、乐于分享知识并回馈他人的数据科学家。Zoumana 是 YouTube 内容创作者，也是 Medium 上的顶尖科技写作者。他热爱演讲、编程与教学。Zoumana 拥有两个硕士学位：其一是在法国巴黎获得的计算机科学硕士，专注于机器学习；其二是在美国德州理工大学获得的数据科学硕士。他的职业生涯始于法国 Groupe OPEN 的软件开发工程师，随后加入 IBM 担任机器学习顾问，为保险公司构建端到端的 AI 解决方案。之后，Zoumana 加入了 Axionable——一家总部位于巴黎与蒙特利尔的首家“可持续 AI”初创公司。在那里，他担任数据科学家，为来自法国、蒙特利尔、新加坡和瑞士的客户交付 AI 产品，主要聚焦 NLP 场景。此外，他将 5% 的时间投入到研发工作。目前，他在 IFC（世界银行集团）担任高级数据科学家。

主题

人工智能

主题

技术