IMDB
一个比较常用的相关模型是：lvwerra/gpt2-imdb
该模型的结构为：
Some weights of GPT2ForSequenceClassification were not initialized from the model checkpoint at lvwerra/gpt2-imdb and are newly initialized: ['score.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
GPT2ForSequenceClassification(
(transformer): GPT2Model(
(wte): Embedding(50257, 768)
(wpe): Embedding(1024, 768)
(drop): Dropout(p=0.1, inplace=False)
(h): ModuleList(
(0-11): 12 x GPT2Block(
(ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
(attn): GPT2SdpaAttention(
(c_attn): Conv1D()
(c_proj): Conv1D()
(attn_dropout): Dropout(p=0.1, inplace=False)
(resid_dropout): Dropout(p=0.1, inplace=False)
)
(ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
(mlp): GPT2MLP(
(c_fc): Conv1D()
(c_proj): Conv1D()
(act): NewGELUActivation()
(dropout): Dropout(p=0.1, inplace=False)
)
)
)
(ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
)
(score): Linear(in_features=768, out_features=2, bias=False)
)
可以看到，该模型为GPT2 + 一个线性层，用于二分类。GPT2在imdb构成的纯文本上经过了微调，而线性层的参数是随机初始化的，需要根据具体任务进行训练。
比如，我们现在想要在此基础上进行reward model的训练，那么我们可以使用如下代码： %accordion% Reward Model
from datasets import load_dataset, load_metric
from transformers import (
AutoModelForSequenceClassification,
AutoTokenizer,
GPT2Tokenizer,
TrainingArguments,
Trainer,
DataCollatorWithPadding,
)
import numpy as np
def tokenize(examples):
outputs = tokenizer(examples["text"], truncation=True, padding=True)
return outputs
def compute_metrics(eval_preds):
metric = load_metric("accuracy")
logits, labels = eval_preds
predictions = np.argmax(logits, axis=-1)
return metric.compute(predictions=predictions, references=labels)
model_name = "lvwerra/gpt2-imdb"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
# tokenizer.padding_side = "left"
model.config.pad_token_id = model.config.eos_token_id
ds = load_dataset("imdb")
tokenized_ds = ds.map(tokenize, batched=True)
print(tokenized_ds["train"][0])
training_args = TrainingArguments(
num_train_epochs=2,
output_dir="gpt2-rm-imdb",
# push_to_hub=True,
per_device_train_batch_size=1,
per_device_eval_batch_size=1,
evaluation_strategy="epoch",
save_strategy="epoch",
)
data_collator = DataCollatorWithPadding(tokenizer)
trainer = Trainer(
model=model,
tokenizer=tokenizer,
data_collator=data_collator,
args=training_args,
train_dataset=tokenized_ds["train"],
eval_dataset=tokenized_ds["test"],
compute_metrics=compute_metrics,
)
trainer.train()
AutoModelForSequenceClassification在 Hugging Face 库中默认会使用交叉熵损失函数（CrossEntropyLoss）。
训练完成之后，logits(pos) - logits(neg)即为reward。