当前位置：首页 > article >正文

基于Python的自然语言处理系列（36）：使用PyTorch微调（无需Trainer）

article 2024/10/18 9:50:09

本篇文章将展示如何通过PyTorch手动实现预训练模型的微调，而不依赖Trainer类。我们将以微软研究院的MRPC（Microsoft Research Paraphrase Corpus）数据集为例。MRPC包含句对和标签，用于判断两句是否为语义相同的重述，非常适合用于探索文本分类任务的微调流程。我们会从数据预处理到训练与评估，逐步演示整个过程，最后介绍如何通过Accelerate库提升多设备上的训练效率。

1. 加载数据集和预处理

首先加载MRPC数据集，并通过BERT的Tokenizer进行分词。我们使用glue任务中的MRPC数据集，并配置BERT的预训练检查点。

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding

# 加载数据集和BERT tokenizer
raw_datasets = load_dataset("glue", "mrpc")
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# 定义分词函数
def tokenize_function(example):
    return tokenizer(example["sentence1"], example["sentence2"], truncation=True)

# 应用分词并生成数据集
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

数据集后处理

我们需要移除不必要的列，并将标签列重命名为模型期望的labels。

tokenized_datasets = tokenized_datasets.remove_columns(["sentence1", "sentence2", "idx"])
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
tokenized_datasets.set_format("torch")

数据加载器

我们创建训练和验证数据的加载器：

from torch.utils.data import DataLoader

train_dataloader = DataLoader(
    tokenized_datasets["train"], shuffle=True, batch_size=8, collate_fn=data_collator
)
eval_dataloader = DataLoader(
    tokenized_datasets["validation"], batch_size=8, collate_fn=data_collator
)

2. 定义模型和优化器

我们使用BERT模型，并通过AdamW优化器和线性学习率调度器配置训练过程：

from transformers import AutoModelForSequenceClassification, get_scheduler
from torch.optim import AdamW

model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=5e-5)

num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
    "linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
)

使用GPU（如果可用）

import torch

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)

3. 训练循环

在训练过程中，我们遍历所有训练数据，并在每次反向传播后更新模型参数。

from tqdm.auto import tqdm

progress_bar = tqdm(range(num_training_steps))
model.train()

for epoch in range(num_epochs):
    for batch in train_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

4. 评估循环

我们使用Evaluate库对验证数据进行评估。

import evaluate

metric = evaluate.load("glue", "mrpc")
model.eval()

for batch in eval_dataloader:
    batch = {k: v.to(device) for k, v in batch.items()}
    with torch.no_grad():
        outputs = model(**batch)

    logits = outputs.logits
    predictions = torch.argmax(logits, dim=-1)
    metric.add_batch(predictions=predictions, references=batch["labels"])

print(metric.compute())

5. 使用Accelerate加速分布式训练

我们可以通过Accelerate库轻松启用多GPU或TPU上的训练。

from accelerate import Accelerator

accelerator = Accelerator()
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=3e-5)

train_dl, eval_dl, model, optimizer = accelerator.prepare(
    train_dataloader, eval_dataloader, model, optimizer
)

progress_bar = tqdm(range(num_training_steps))
model.train()

for epoch in range(num_epochs):
    for batch in train_dl:
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)