微调HuggingFace中提供的模型做图像分类

1,466 阅读8分钟

环境配置

  • 安装依赖库
pip install transformers datasets evaluate

微调模型

数据预处理

下面以resnet-50为例

from transformers import AutoImageProcessor
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor,RandomHorizontalFlip

# 定义转换函数
def transforms(examples):
    examples["pixel_values"] = [_transforms(img.convert("RGB")) for img in examples["image"]]
    del examples["image"]
    return examples

def get_dataset(train_dir,test_dir):
  # 加载本地数据集
  dataset = load_dataset("imagefolder",data_dir=train_dir)
  ds_val = load_dataset("imagefolder",data_dir=test_dir)

  # 整合数据标签和下标
  dataset["test"] = ds_val["train"]
  print("dataset:",dataset)
  return dataset

# 模型名
checkpoint = "microsoft/resnet-50"
# 类似于Pytorch的ImageFloder方式,每个类的数据存放在一个文件夹下
train_dir = "your_train_dir" 
test_dir = "your_test_dir"

# 预处理
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),RandomHorizontalFlip(), ToTensor(), normalize])

# 获取数据集
dataset = get_dataset(train_dir=train_dir,test_dir=test_dir)
ds = dataset.with_transform(transforms)
# label和id相互映射字典
labels = dataset["train"].features["label"].names
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
    label2id[label] = i
    id2label[i] = label

# 用于从train_dataset或eval_dataset的元素列表形成批处理的函数
data_collator = DefaultDataCollator()

  • 使用ImageFolder方式加载训练集和验证集,训练集数据类似于Pytorch的方式准备即可,一个类的数据保存在一个文件夹中

准备模型

# 定义模型
model = AutoModelForImageClassification.from_pretrained(
    checkpoint,
    num_labels=len(labels),
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes = True

)
  • 如果是微调模型,ignore_mismatched_sizes一定要设置成True,不然会报下面的错误,提示模型的维度不匹配。在某些情况下,例如当您试图加载一个包含更多类别的预训练模型时,可能会遇到权重尺寸不匹配的问题。ignore_mismatched_sizes 参数可以帮助您解决这个问题,并避免因权重大小不匹配而导致的错误。
RuntimeError: Error(s) in loading state_dict for ResNetForImageClassification:
	size mismatch for classifier.1.weight: copying a param with shape torch.Size([1000, 2048]) from checkpoint, the shape in current model is torch.Size([5, 2048]).
	size mismatch for classifier.1.bias: copying a param with shape torch.Size([1000]) from checkpoint, the shape in current model is torch.Size([5]).
	You may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.
  • 如果是想使用迁移学习的方式来训练,也就是只想训练最后的全连接层参数,可以使用如下代码
# 定义模型
model = AutoModelForImageClassification.from_pretrained(
    checkpoint,
    num_labels=num_labels,
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes = True
)

# 替换最后一层全连接层
in_features = model.classifier[1].in_features
model.classifier[1] = torch.nn.Linear(in_features, num_labels)

# 冻结前面几层
for param in model.parameters():
  param.requires_grad = False
for param in model.classifier.parameters():
  param.requires_grad = True

total_trainable_params = sum(
    p.numel() for p in model.parameters() if p.requires_grad)
print(f'可训练参数总数:{total_trainable_params}')   

构建训练器

# 计算指标
def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)
    
# 训练超参数
training_args = TrainingArguments(
    output_dir="hf/resnet",
    remove_unused_columns=False,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-4,
    save_total_limit = 3,
    per_device_train_batch_size=32,
    gradient_accumulation_steps=1,
    per_device_eval_batch_size=32,
    num_train_epochs=20,
    warmup_ratio=0.1,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=False,
)

# 训练参数
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=ds["train"],
    eval_dataset=ds["test"],
    tokenizer=image_processor,
    compute_metrics=compute_metrics,
)

如果在jupyter中训练会打印如下的训练过程

  • 设置一些超参数、训练参数等

以下是 TrainingArguments 类的一些常用参数:

  • output_dir:输出目录,用于保存模型检查点和训练日志。
  • num_train_epochs:进行训练的总轮次(epochs)。
  • per_device_train_batch_size:每个设备上的训练批次大小。
  • per_device_eval_batch_size:每个设备上的评估批次大小。
  • evaluation_strategy:评估策略(例如,“steps”或“epoch”),确定在训练过程中何时评估模型。
  • logging_dir:用于存储训练日志文件的目录。
  • seed:为了确保实验的可重复性,设置随机种子。
  • learning_rate:模型的学习率。
  • save_strategy:模型检查点的保存策略(例如,“steps”或“epoch”),确定何时保存模型。
  • logging_steps:记录训练指标的时间间隔(以步数计)。
  • save_steps:保存模型检查点的时间间隔(以步数计)。
  • save_total_limit:存储在输出目录中的最大模型检查点数量。

TrainingArguments详细的参数列表

  • 其中的epochs和step的关系如下:
total_steps = num_epochs * (num_train_samples / batch_size)

比如训练样本是1356,每批次大小是32,设置10个epochs,那总的迭代步数如下:

total_steps = 10 * (1356 / 32) ≈ 10 * 42.375 = 423.75 ~=424

以下是 Trainer 类的一些常用参数:

  • model:您要训练或评估的预训练模型。该参数必须是一个从 torch.nn.Module 继承而来的模型实例。
  • args:一个 TrainingArguments 对象,其中包含训练相关的设置。诸如学习率,批次大小,训练轮数,输出目录等。
  • train_dataset:用于训练模型的数据集。这必须是一个 datasets.Dataset 实例或支持实现 __len____getitem__ 的任何自定义数据集。
  • eval_dataset:用于评估模型的数据集。这与 train_dataset 类似,但用于在训练过程中评估模型的性能。
  • tokenizer:用于预处理文本输入的分词器。它通常是一个从 PreTrainedTokenizer 类派生出来的实例,与要训练的预训练模型相对应。
  • compute_metrics:一个可选参数,用于计算特定任务的评估指标。该参数是一个接收 EvalPrediction(包含模型预测和实际目标值)的函数,并返回一个指标(如准确性、F1 分数等)的字典。这有助于在训练过程中跟踪模型的性能。

Trainer详细的参数列表

开始训练

# 开始训练
train_results = trainer.train()
# 保存模型
trainer.save_model()
trainer.log_metrics("train", train_results.metrics)
trainer.save_metrics("train", train_results.metrics)
trainer.save_state()

上面第一行中的

[ 2254/2500 21:23 < 02:20, 1.75 it/s, Epoch 45.06/50]

的含义如下:

[<完成的步骤>/<总的步骤> <已经过的时间> <预计剩余时间>, <每秒迭代数>, Epoch <完成的纪元>/<总纪元>]

比如上面的含义如下

  • 2254/2500:已完成 2254 步,总共计划进行 2500 步。
  • 21:23:从训练开始已经过去 21分23秒。
  • 02:20:预计还需要 2分20秒才能完成整个训练过程。
  • 1.75 it/s:平均每秒钟处理 1.75 个迭代(分批次)。
  • Epoch 45.06/50:已经完成了 45.06 个 epochs,计划进行 50 个 epochs。这里的45.06 表示模型已经完成了45轮 epoch,并开始进行下一轮的 0.06 个 epochs。

这些统计数据有助于了解训练过程的进度、速度以及剩余时间。不过需要注意的是,预计剩余的时间可能会随着计算资源负载的变化而发生变化。

开始评估

metrics = trainer.evaluate()
trainer.log_metrics("eval",metrics)
trainer.save_metrics("eval", metrics)

打印出来如下结果


***** eval metrics *****
  epoch                   =       50.0
  eval_accuracy           =     0.8736
  eval_loss               =     0.4681
  eval_runtime            = 0:00:14.76
  eval_samples_per_second =     11.784
  eval_steps_per_second   =      0.745

推理

使用pipline方式

# 使用pipline
from transformers import pipeline

test_dir = 'your_test_dir'  # 替换为你的数据目录
ds = load_dataset("imagefolder", data_dir=test_dir)
print("ds:",ds)
image = ds["train"]["image"][1]

classifier = pipeline("image-classification", model="./hf/resnet/")
classifier(image)

手动加载模型

# 手动方式
from transformers import AutoImageProcessor
import torch
from transformers import AutoModelForImageClassification
from PIL import Image

test_dir = 'your_test_dir'  # 替换为你的数据目录
ds = load_dataset("imagefolder", data_dir=test_dir)
print("ds:",ds)
image = ds["train"]["image"][1]

# 指定本地模型和分词器的路径
model_path = "./hf/resnet/"
tokenizer_path = "./hf/resnet/"

# 需要转换成三通道格式图片,因为AutoImageProcessor需要一个RGB(三通道)格式的图像。
image1 = image.convert("RGB")
image_processor = AutoImageProcessor.from_pretrained(tokenizer_path)
inputs = image_processor(image1, return_tensors="pt")


model = AutoModelForImageClassification.from_pretrained(model_path)
with torch.no_grad():
    logits = model(**inputs).logits
predicted_label = logits.argmax(-1).item()
model.config.id2label[predicted_label]

批量验证模型的识别率

from transformers import ResNetForImageClassification
from torchvision import transforms, datasets
import torch

model_path = "./hf/resnet/"

# 加载模型
model = ResNetForImageClassification.from_pretrained(model_path)
model.eval()  # 设置模型为评估模式

# 定义预处理步骤
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 加载数据
test_dir = 'your_test_dir'  # 替换为你的数据目录
dataset = datasets.ImageFolder(test_dir, transform=preprocess)

# 创建数据加载器
batch_size = 32  # 设置批量大小
data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False)

# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 运行模型进行预测并计算损失
total_loss = 0.0

# 运行模型进行预测
with torch.no_grad():
    all_preds = []
    for batch in data_loader:
        inputs, _ = batch
        outputs = model(inputs)
        # transformers库中,模型的输出是一个特定的对象,而不是一个张量。对于图像分类模型,
        _, preds = torch.max(outputs.logits, 1)
        loss = criterion(outputs.logits, labels)  # 计算损失
        all_preds.extend(preds.cpu().numpy())
        total_loss += loss.item()


# 获取预测结果
all_preds = torch.tensor(all_preds)
labels = dataset.targets

# 计算识别率
accuracy = (all_preds == torch.tensor(labels)).sum().item() / len(labels)
# 计算平均损失
avg_loss = total_loss / len(data_loader)

print(f'整体测试集上的Loss: {avg_loss:.5f}')
print(f'整体数据集上的准确率Acc: {accuracy * 100:.5f}%')


如果想可视化识别的结果,可以如下操作:

# 可视化训练图片,也可以使用TensorBoard
import matplotlib.pyplot as plt
import numpy as np
import torchvision

# 展示图片
def imshow(img):
    img = img/2 +0.5 #  对图像进行逆归一化。在进行训练前,图像很可能已经被归一化,其像素值范围在[-1, 1]之间。逆归一化之后,像素值范围将恢复到 [0, 1],从而能够以原始形式正确显示图像。
    npimg = img.numpy() # pytorch张量转换成numpy
    plt.imshow(np.transpose(npimg,(1,2,0))) # np.transpose() 对 NumPy 数组 npimg 进行转置,第1个轴(高度 H)变为新的第 0 个轴 2=>1 0=>2
    plt.show()

def printLable(label_list,nrow):
    output = ""
    for i, item in enumerate(label_list):
      output += str(item) + "\t"
      if (i + 1) % nrow == 0:
          output += "\n"
    print( output)

# 可视化(数据集比较少时)
classes = model.config.id2label
result = (all_preds == torch.tensor(labels)).cpu().numpy()

nrow = 10 # 每行显示的图片数量

print(images[0].shape)
# 展示图片
batch_images = torch.stack(images)
print("数据大小:",batch_images.shape)
temp = torchvision.utils.make_grid(batch_images,nrow=nrow, padding=5, pad_value=1)
imshow(temp)

# 使用列表推导式,将 label_list 转换为对应的 class 标签列表
class_labels = [classes[i] for i in all_preds.cpu().numpy()]
print("------------------预测类型--------------------------")

printLable(class_labels,nrow)
print("------------------正确性--------------------------")
printLable(result,nrow)

附录

  • 完整的训练代码如下
from datasets import load_dataset
import evaluate
from transformers import AutoImageProcessor
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor,RandomHorizontalFlip
import numpy as np
from transformers import AutoModelForImageClassification, TrainingArguments, Trainer
from transformers import DefaultDataCollator

# 定义转换函数
def transforms(examples):
    examples["pixel_values"] = [_transforms(img.convert("RGB")) for img in examples["image"]]
    del examples["image"]
    return examples

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)

def get_dataset(train_dir,test_dir):
  # 加载本地数据集
  dataset = load_dataset("imagefolder",data_dir=train_dir)
  ds_val = load_dataset("imagefolder",data_dir=test_dir)

  # 整合数据标签和下标
  dataset["test"] = ds_val["train"]
  print("dataset:",dataset)
  return dataset

# 模型名
checkpoint = "microsoft/resnet-50"
# 类似于Pytorch的ImageFloder方式,每个类的数据存放在一个文件夹下
train_dir = "your_train_dir" 
test_dir = "your_test_dir"

# 预处理
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),RandomHorizontalFlip(), ToTensor(), normalize])

# 获取数据集
dataset = get_dataset(train_dir=train_dir,test_dir=test_dir)
ds = dataset.with_transform(transforms)
# label和id相互映射字典
labels = dataset["train"].features["label"].names
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
    label2id[label] = i
    id2label[i] = label

# 用于从train_dataset或eval_dataset的元素列表形成批处理的函数
data_collator = DefaultDataCollator()


# 定义模型
model = AutoModelForImageClassification.from_pretrained(
    checkpoint,
    num_labels=len(labels),
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes = True

)

# 训练超参数
training_args = TrainingArguments(
    output_dir="hf/resnet",
    remove_unused_columns=False,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-4,
    save_total_limit = 3,
    per_device_train_batch_size=32,
    gradient_accumulation_steps=1,
    per_device_eval_batch_size=32,
    num_train_epochs=20,
    warmup_ratio=0.1,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=False,
)

# 训练参数
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=ds["train"],
    eval_dataset=ds["test"],
    tokenizer=image_processor,
    compute_metrics=compute_metrics,
)

# 开始训练
train_results = trainer.train()
print("train_results:",train_results)
# 保存模型
trainer.save_model()
trainer.log_metrics("train", train_results.metrics)
trainer.save_metrics("train", train_results.metrics)
trainer.save_state()

metrics = trainer.evaluate()
trainer.log_metrics("eval",metrics)
trainer.save_metrics("eval", metrics)

参考

微调Hugging Face中图像分类模型

Image classification

Load image data

huggingface官网关于"imagefolder"的说明