LoRA微调BERT情感分析

用LoRA微调BERT进行情感分析

pip install torch transformers datasets peft accelerate #环境准备

• 数据选用IMDb电影评论数据集

from datasets import load_dataset

dataset = load_dataset("imdb")
print(dataset["train"][0])  

• 加载基础模型与Tokenizer

from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# 原始参数量统计
total_params = sum(p.numel() for p in model.parameters())
print(f"原始模型参数总量：{total_params:,}")  

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=8,  # 秩选择8，经测，秩在16收益仅仅增加0.2%，说明收益拐点就在8附近
    lora_alpha=16,  # 缩放系数
    target_modules=["query", "value"], # Query和Value层直接参与语义特征交互
    lora_dropout=0.1,   #防止过拟合
    bias="none",  # 不训练偏置项
    modules_to_save=["classifier"]  # 预训练模型的原始分类头是随机初始化的，必须参与训练
)
# 注入
peft_model = get_peft_model(model, lora_config)

# 可训练参数量统计
trainable_params = sum(p.numel() for p in peft_model.parameters() if p.requires_grad)
print(f"可训练参数占比：{trainable_params}/{total_params} ({100*trainable_params/total_params:.2f}%)") 
# 约0.8%参数被训练（约0.9M）

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./lora_imdb", #输出到./lora_imdb
    learning_rate=3e-4,        # 3倍学习率用于突破局部最优
    per_device_train_batch_size=8, 
    num_train_epochs=3,
    logging_steps=100,
    save_strategy="epoch",
    fp16=True                  # 启用混合精度加速
)

def tokenize_function(examples): # 分词
    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=256)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

trainer = Trainer(
    model=peft_model,
    args=training_args,
    train_dataset=tokenized_datasets["train"].shuffle().select(range(1000)),  # 演示用子集
    eval_dataset=tokenized_datasets["test"].shuffle().select(range(200)),
    tokenizer=tokenizer,
)

trainer.train()
peft_model.save_pretrained("lora_imdb_output")  # 仅保存LoRA权重

from peft import PeftModel
base_model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2) # 加载基础模型
model = PeftModel.from_pretrained(base_model, "lora_imdb_output") # 注入LoRA
model = model.merge_and_unload() # 合并权重加速
inputs = tokenizer("This movie was fantastic!", return_tensors="pt")# 预测示例
outputs = model(**inputs)
pred = outputs.logits.argmax().item()
print("正面" if pred == 1 else "负面")  # 输出：正面