BERT 微调 · 中文分类最小可行版 | 小朵科技工作室 博客

前面聊过 BERT，但都是直接用预训练模型。这篇讲怎么用自己的数据微调。

最常见的任务：评论分类（好评/差评/中性）。

最小可行流程：

1. 1

text,label
"这个手机真不错",1
"垃圾",0

1. 2

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(
    texts, labels, test_size=0.2
)

1. 3

from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments

tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertForSequenceClassification.from_pretrained(
    'bert-base-chinese', num_labels=3
)

def tokenize(batch):
    return tokenizer(batch['text'], padding=True, truncation=True, max_length=128)

train_dataset = Dataset.from_dict({'text': X_train, 'label': y_train}).map(tokenize, batched=True)
test_dataset = Dataset.from_dict({'text': X_test, 'label': y_test}).map(tokenize, batched=True)

training_args = TrainingArguments(output_dir='./results', num_train_epochs=3, per_device_train_batch_size=16)
trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset, eval_dataset=test_dataset)
trainer.train()

硬件要求：有 GPU 推荐，没有也能跑（慢 10 倍）。

数据量经验值：< 1000 条效果一般，> 5000 条开始起飞，> 10000 条基本能压过商业 API。

---

本文由 BBZ · 小朵科技工作室 出品。