一、任务背景
最近几年,大模型军备竞赛愈演愈烈。而对我们普通人而言,更多的时候只是把大模型当成AI助手来使用,例如最近大火的DeepSeek。受限于算力资源以及数据规模,普通个体甚至是没有服务器集群的小公司,想定制化大模型变成了极其困难的事情。然而,随着大模型微调技术的进一步发展,我们可以仅使用少量算力来完成一次个性化的大模型微调,这可是给小团队带来了巨大的便利性。
本文基于HuggingFace教程中的PEFT教程,进行大模型的微调实战。目前,高效训练大模型的主流方法之一是在如注意力模块中插入一个较小的可训练矩阵,作为在微调期间要学习的增量权重矩阵的低秩分解。预训练模型的原始权重矩阵被冻结,在训练期间仅更新较小的矩阵。这减少了可训练参数的数量,减少了内存使用量和训练时间,大大降低了大模型微调的成本。有几种不同的方法可以将权重矩阵表示为低秩分解,但低秩适应(LoRA)是最常见的方法。当然,HuggingFace开发的PEFT 库也支持例如LoHa、LoKr和AdaLoRA等其他几种LoRA变体。此外,PEFT 支持X-LoRA低秩专家混合方法。 在这篇博客中,我们将跟随着PEFT官方教程学习如何使用低秩分解方法快速训练一个图像分类模型,用于识别图像中所示食物的类别。
二、python实战
1、数据准备
教程中使用的是Food-101数据集,该数据集包含 101 种食物类别的图像。每个食物类别都用一个整数进行标记,我们将创建一个label2id和id2label字典,以便将整数映射到其类别标签。
from datasets import load_dataset
ds = load_dataset("food101")
labels = ds["train"].features["label"].names
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
label2id[label] = i
id2label[i] = label
接着,我们加载一个图像处理器以正确调整训练和评估图像的大小,并归一化像素值。
from transformers import AutoImageProcessor
from torchvision.transforms import (
CenterCrop,
Compose,
Normalize,
RandomHorizontalFlip,
RandomResizedCrop,
Resize,
ToTensor,
)
image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
train_transforms = Compose(
[
RandomResizedCrop(image_processor.size["height"]),
RandomHorizontalFlip(),
ToTensor(),
normalize,
]
)
val_transforms = Compose(
[
Resize(image_processor.size["height"]),
CenterCrop(image_processor.size["height"]),
ToTensor(),
normalize,
]
)
def preprocess_train(example_batch):
example_batch["pixel_values"] = [train_transforms(image.convert("RGB")) for image in example_batch["image"]]
return example_batch
def preprocess_val(example_batch):
example_batch["pixel_values"] = [val_transforms(image.convert("RGB")) for image in example_batch["image"]]
return example_batch
下面,对训练集和测试集分别应用变换函数。同时,创建一个数据整理器用于训练和测试。
import torch
train_ds = ds["train"]
val_ds = ds["validation"]
train_ds.set_transform(preprocess_train)
val_ds.set_transform(preprocess_val)
def collate_fn(examples):
pixel_values = torch.stack([example["pixel_values"] for example in examples])
labels = torch.tensor([example["label"] for example in examples])
return {"pixel_values": pixel_values, "labels": labels}
2、模型准备
首先,加载一个基础模型,这里使用的是google/vit-base-patch16-224-in21k模型。接着,将label2id和id2label字典传递给模型,以便它知道如何将整数标签映射到它们的类别标签。
from transformers import AutoModelForImageClassification, TrainingArguments, Trainer
model = AutoModelForImageClassification.from_pretrained(
"google/vit-base-patch16-224-in21k",
label2id=label2id,
id2label=id2label,
ignore_mismatched_sizes=True,
)
3、LoRA配置
低秩自适应(LoRA)将权重更新矩阵分解为两个较小的矩阵。这些低秩矩阵的大小由其秩或r决定。较高的秩意味着模型有更多的参数需要训练,但也意味着模型有更大的学习能力。我们还需要指定target_modules,它决定了较小的矩阵插入的位置。这里,我们将针对注意力块的查询和值矩阵。参数设置如下:
- r:秩,即rank;
- lora_alpha:用于 LoRA 缩放;
- target_modules:要应用适配器的模块名称,如果指定了这个参数,只有具有指定名称的模块才会被替换(传入一个字符串时,将进行正则表达式匹配;当传入一个字符串列表时,将进行精确匹配,或者检查模块名称是否以任何传入的字符串结尾;如果这个参数指定为“all-linear”,那么所有线性/Conv1D 模块都会被选中,但输出层除外;如果没有指定这个参数,将根据模型架构选择模块);
- bias偏差:可以是“none”、“all”或“lora_only”。如果是“all”或“lora_only”,相应的偏差将在训练期间更新。这意味着即使禁用adapter,模型也不会产生与未进行适配的基础模型相同的输出。
- modules_to_save:除了adapter层之外的模块列表,这些模块将被设置为可训练的,并保存在最终的checkpoint中。
from peft import LoraConfig, get_peft_model
config = LoraConfig(
r=16,
lora_alpha=16,
target_modules=["query", "value"],
lora_dropout=0.1,
bias="none",
modules_to_save=["classifier"],
)
model = get_peft_model(model, config)
model.print_trainable_parameters()
4、模型训练
对于训练,我们使用来自 Transformers 的Trainer类。Trainer包含一个 PyTorch 训练循环。下面看看一些参数设定:
- output_dir:输出路径。
- remove_unused_columns:可选参数,默认为True,表示是否自动移除模型前向传播方法未使用的列。
- eval_strategy:训练过程中的评估策略,这里是每个epoch评估一次,也可以设置为固定steps评估一次或者不评估。
- save_strategy:训练权重保存策略,这里是每个epoch保存一次。
- learning_rate:学习率。
- per_device_train_batch_size:每个设备上的训练batch_szie。
- gradient_accumulation_steps:在执行反向传播/更新步骤之前,用于累积梯度的更新步数。
- per_device_eval_batch_size:每个设备上的验证batch_size。
- fp16:是否使用16位(混合)精度训练而非32位训练。
from transformers import TrainingArguments, Trainer
peft_model_id = f"model/vit-base-patch16-224-in21k-lora"
batch_size = 128
args = TrainingArguments(
output_dir=peft_model_id,
remove_unused_columns=False,
eval_strategy="epoch",
save_strategy="epoch",
learning_rate=5e-3,
per_device_train_batch_size=batch_size,
gradient_accumulation_steps=4,
per_device_eval_batch_size=batch_size,
fp16=True,
num_train_epochs=5,
logging_steps=10,
load_best_model_at_end=True,
label_names=["labels"],
)
训练模型。
trainer = Trainer(
model,
args,
train_dataset=train_ds,
eval_dataset=val_ds,
tokenizer=image_processor,
data_collator=collate_fn,
)
trainer.train()
5、模型应用
from peft import PeftConfig, PeftModel
from transformers import AutoImageProcessor
from PIL import Image
import requests
config = PeftConfig.from_pretrained("model/vit-base-patch16-224-in21k-lora")
model = AutoModelForImageClassification.from_pretrained(
config.base_model_name_or_path,
label2id=label2id,
id2label=id2label,
ignore_mismatched_sizes=True,
)
model = PeftModel.from_pretrained(model, "model/vit-base-patch16-224-in21k-lora")
url = "https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/beignets.jpeg"
image = Image.open(requests.get(url, stream=True).raw)
encoding = image_processor(image.convert("RGB"), return_tensors="pt")
with torch.no_grad():
outputs = model(**encoding)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])