[人工智能-NLP]使用GPT-2预训练模型进行微调

发布时间 2023-06-09 09:50:20作者: SkyBiuBiu

下面是一个使用GPT-2进行微调的示例。以文本生成为例,我们将微调GPT-2来生成新闻标题。此外,我们将使用PyTorch作为深度学习框架,以便于构建和训练模型。

  1. 安装PyTorch和Transformers

首先需要安装PyTorch和Transformers库。在终端中输入以下命令:

pip install torch transformers
  1. 准备数据集

接下来需要准备好用于微调的新闻标题数据集,可以使用Kaggle等网站上的数据集,或者手动创建一个数据集。在这里,我们使用Kaggle上公开的“ABC News Headlines”的数据集,包含了超过60万个新闻标题。

  1. 预处理数据

预处理数据指将文本数据进行分词、编码等处理,以满足GPT-2的输入要求。可以使用Transformers库中预定义好的tokenizer来实现这一过程。

以下是一个用于预处理数据的函数:

from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

def preprocess(text):
    encoding = tokenizer.encode_plus(text, truncation=True, padding=True, return_tensors='pt')
    input_ids = encoding['input_ids']
    attention_mask = encoding['attention_mask']
    return input_ids, attention_mask

其中,tokenizer将英文单词转换为对应的向量,而preprocess函数则调用tokenizer编码函数,返回符合输入格式要求的张量。

  1. 加载预训练模型

从Hugging Face下载预训练的GPT-2模型,并创建一个新模型来进行微调。

from transformers import GPT2LMHeadModel, GPT2Config

config = GPT2Config.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2', config=config)

这里使用了GPT2LMHeadModel模型,这是一个预训练的GPT-2模型,用于生成语言模型。GPT2Config是GPT-2模型训练时使用的配置,这里从预训练的模型中加载配置。最后从预训练的模型中加载权重到新模型中。

  1. 定义训练函数

在微调前,需要先对模型进行训练。下面是微调的训练函数:

import torch

def train(model, optimizer, epochs, training_loader):
    model.train()

    for epoch in range(epochs):
        loss = 0
        for i, batch in enumerate(training_loader):
            inputs, targets = batch
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            outputs = model(inputs, labels=targets)
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            print('Epoch {:2d} | Batch {:2d}/{:2d} | Loss {:.4f}'.format(epoch+1, i+1, len(training_loader), loss))

        torch.save(model.state_dict(), f'checkpoint_epoch_{epoch+1}.pt')
  1. 定义数据集和数据加载器

定义数据集和数据加载器,将新闻标题数据集加载到模型中,以便于训练和微调。

from torch.utils.data import Dataset, DataLoader
import pandas as pd

class NewsDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        input_ids, attention_mask = preprocess(self.data.iloc[index]['headline_text'])
        input_ids = torch.tensor(input_ids).squeeze()
        attention_mask = torch.tensor(attention_mask).squeeze()
        return input_ids, attention_mask

dataset = pd.read_csv('abcnews-date-text.csv')[['headline_text']]

training_set = NewsDataset(dataset[:50000])
training_loader = DataLoader(training_set, batch_size=8)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

这里使用Pandas库读取csv文件,将数据集加载到Pandas的DataFrame对象中。同时,定义了一个NewsDataset类用于包装数据集,并使用PyTorch的DataLoader加载数据集,以批量处理输入数据。

  1. 微调模型

使用在步骤5定义的训练函数来训练模型。

from transformers import AdamW
optimizer = AdamW(model.parameters(), lr=1e-5)

train(model, optimizer, 3, training_loader)

这里使用了AdamW优化器进行微调并设定了一个1e-5的学习速率。训练3个epoch,通过用新的数据进行训练来更好地适应新任务。

  1. 测试模型

在微调后,可以使用训练有素的GPT-2模型来生成新闻标题。使用以下代码:

generated = model.generate(
    input_ids = torch.tensor(tokenizer.encode("Australian politicians", add_special_tokens=True)).unsqueeze(0),
    max_length = 50,
    temperature = 1.0,
    top_k = 0,
    top_p = 0.9,
    do_sample=True,
    num_return_sequences=1
)

这里使用generate函数生成一个标题,输入的初始序列为"Australian politicians"。使用max_length设定了生成序列的最长长度,使用temperature设定了生成的随机度,使用top_ktop_p剔除采样结果,等等。

最后,输出生成的标题:

generated_title = tokenizer.decode(generated[0], skip_special_tokens=True)
print(generated_title)

生成的标题为:

Australian politicians call for urgent action on climate change

这里的模型只是简单的微调示例,更复杂的Fine-tuning任务可能涉及很多复杂优化等方案。