【Pytorch和Keras】使用transformer库进行图像分类

现在大多数的模型都会上传到huggface平台进行统一的管理，transformer库能关联到huggface中对应的模型，并且提供简洁的transformer模型调用，这大大提高了开发人员的开发效率。本博客主要利用transformer库实现一个简单的模型微调，以进行图像分类的任务。

一、环境准备

使用终端命令行安装对应的第三方包，具体安装命令输入如下：

pip install transformers datasets evaluate

二、基于Pytorch的预训练模型

由于下面这些内容需要在huggface上申请账号权限，才能进行模型和数据集加载，如果之前有从huggface上拉取模型和数据集的经验，可以略过，如果没有配置过，可以参考笔者之前的文章https://blog.csdn.net/qq_40734883/article/details/143922095，然后直接申请Write权限就可以。

在这里插入图片描述
后续所有涉及到的数据集food101和transformer模型都需要参考上述文章进行直接下载，才能运行整个程序，或者在google的colab直接运行。

如果在google的colab上运行，请提前设置好电脑的GPU资源，同时加入huggface登录代码，具体如下：

from huggingface_hub import notebook_login
notebook_login()

运行之后会提示进行token输入，按之前获取到的token输入即可。

1、准备数据集

这里以food101数据集作为微调数据集，在imagenet-21k上训练完成的transformer模型（vit-base-patch16-224）进行优化

from datasets import load_dataset

food = load_dataset("food101", split="train[:5000]")


# 划分数据集，训练集:测试集=8:2，food有两个键：一个train，一个test
food = food.train_test_split(test_size=0.2)  

# 标签转换
labels = food["train"].features["label"].names
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
    label2id[label] = str(i)
    id2label[str(i)] = label

id2label为通过id访问标签的字典，后续会使用到。

2、加载预训练模型

from transformers import AutoImageProcessor

checkpoint = "google/vit-base-patch16-224-in21k"   # ImageNet-21k上的预训练模型
image_processor = AutoImageProcessor.from_pretrained(checkpoint)  # 从huggface拉取并加载模型

3、使用pytorch进行模型构建

from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor

# 数据预处理操作定义
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])

# 对原始数据进行RGB及字典化
def transforms(examples):
    examples["pixel_values"] = [_transforms(img.convert("RGB")) for img in examples["image"]]
    del examples["image"]
    return examples

food = food.with_transform(transforms)

# 验证
import evaluate
# 指定验证过程中的评价指标-准确率
accuracy = evaluate.load("accuracy")

import numpy as np
def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

训练设置和运行，具体输入代码如下：

# 整合训练中的数据，以便在模型训练或评估过程中使用
from transformers import DefaultDataCollator
data_collator = DefaultDataCollator()

from transformers import AutoModelForImageClassification, TrainingArguments, Trainer

# 初始化模型
model = AutoModelForImageClassification.from_pretrained(
    checkpoint,
    num_labels=len(labels),
    id2label=id2label,
    label2id=label2id,
)

# 设置模型优化参数
training_args = TrainingArguments(
    output_dir="my_awesome_food_model",
    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",
    push_to_hub=True,
)

# 初始化训练实例
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=food["train"],
    eval_dataset=food["test"],
    tokenizer=image_processor,
    compute_metrics=compute_metrics,
)

trainer.train()  # 开始训练

trainer.push_to_hub()  # 推送到huggfacehub

经过上述设置训练完成之后，会将模型微调结果推送到huggface平台，如果不想推送，可以不运行相关的命令行，并且training_args中的push_to_hub=False。

训练结果如下图所示：

在这里插入图片描述
默认需要选择是否关联wandb，如果不想选择，直接根据设置提示跳过即可。

如果选择了推送到huggfacehub（trainer.push_to_hub() ）的话，在个人的huggface上会有一个名为my_awesome_food_model的模型，里面包含了模型训练的各个参数设置和测试结果。

在这里插入图片描述

三、基于keras的预训练模型

使用transflow的keras API 进行模型的搭建，具体代码如下：

from transformers import create_optimizer

# 超参数设置
batch_size = 16
num_epochs = 5
num_train_steps = len(food["train"]) * num_epochs
learning_rate = 3e-5
weight_decay_rate = 0.01

# 定义优化方式和策略
optimizer, lr_schedule = create_optimizer(init_lr=learning_rate, 
										  num_train_steps=num_train_steps, 
										  weight_decay_rate=weight_decay_rate, 
										  num_warmup_steps=0)

# 定义分类器
from transformers import TFAutoModelForImageClassification
model = TFAutoModelForImageClassification.from_pretrained(checkpoint, id2label=id2label, label2id=label2id)


# converting our train dataset to tf.data.Dataset
tf_train_dataset = food["train"].to_tf_dataset(columns="pixel_values", label_cols="label", shuffle=True, batch_size=batch_size, collate_fn=data_collator)

# converting our test dataset to tf.data.Dataset
tf_eval_dataset = food["test"].to_tf_dataset(columns="pixel_values", label_cols="label", shuffle=False, batch_size=batch_size, collate_fn=data_collator)

# 定义损失函数
from tensorflow.keras.losses import SparseCategoricalCrossentropy
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

model.compile(optimizer=optimizer, loss=loss)


from transformers.keras_callbacks import KerasMetricCallback, PushToHubCallback
# 定义验证指标
metric_callback = KerasMetricCallback(metric_fn=compute_metrics, eval_dataset=tf_eval_dataset)
# 推送到huggface回调函数
push_to_hub_callback = PushToHubCallback(output_dir="food_classifier", 
                                         tokenizer=image_processor, 
                                         save_strategy="no")
callbacks = [metric_callback, push_to_hub_callback]

# 开始训练
model.fit(tf_train_dataset, validation_data=tf_eval_dataset, epochs=num_epochs, callbacks=callbacks)

四、模型测试

这里使用微调好的模型在food101上找一张验证图像进行简单的验证测试，具体代码如下：

# 验证food中验证集的某一张图像
ds = load_dataset("food101", split="validation[-5:-1]")
image = ds["image"][-1]

# visualize image
import matplotlib.pyplot as plt
plt.imshow(image)
plt.axis('off')  
plt.show()

测试图像如下所示：
在这里插入图片描述

from transformers import pipeline
# initialize classifier instance
classifier = pipeline("image-classification", model="my_awesome_food_model")
classifier(image)

from transformers import AutoImageProcessor
import torch
# load pre-trained image processor
image_processor = AutoImageProcessor.from_pretrained("my_awesome_food_model")
inputs = image_processor(image, return_tensors="pt")


from transformers import AutoModelForImageClassification
# laod pre-trained model
model = AutoModelForImageClassification.from_pretrained("my_awesome_food_model")
with torch.no_grad():
    logits = model(**inputs).logits

# 输出测试结果
predicted_label = logits.argmax(-1).item()
print(model.config.id2label[predicted_label])

输出结果如下所示：

Device set to use cuda:0
[{'label': 'ramen', 'score': 0.9517934918403625},
 {'label': 'bruschetta', 'score': 0.7566707730293274},
 {'label': 'hamburger', 'score': 0.7004948854446411},
 {'label': 'chicken_wings', 'score': 0.6275856494903564},
 {'label': 'prime_rib', 'score': 0.5991673469543457}]

预测结果为：ramen。

释义：“ramen”一词源于日语“ラーメン”，是“拉面”的意思。它进一步追溯至汉语“拉面”，是一种起源于中国、流行于日本及其他东亚地区的面条食品。在日本，拉面通常由小麦面粉制成的面条，搭配肉汤和各种配料，如叉烧、鸡蛋、蔬菜等。

五、参考

[1] https://huggingface.co/docs/transformers/main/tasks/image_classification

[2] https://github.com/huggingface/transformers/blob/main/docs/source/en/installation.md