BERT 模型训练实践
平台：Windows 11 + NVIDIA RTX 4060 + CUDA 12.7 + Miniconda + PyTorch + Hugging Face Transformers
使用 Hugging Face 的 Transformers 库，基于 bert-base-uncased 模型进行微调，完成一个句子评分/分类任务的训练与预测流程，并使用 GPU 加速训练。
完整代码地址：https://github.com/Doge2077/learn-bert
技术介绍
BERT 核心思想
BERT（Bidirectional Encoder Representations from Transformers）是由Google在2018年提出的预训练语言模型，基于Transformer架构，通过大规模无监督语料训练，能够捕捉文本的深层语义和上下文信息。以下是BERT的核心原理及其各层功能的详细解析：
一、BERT的核心原理
- 双向上下文建模：
- BERT通过Transformer的自注意力机制（Self-Attention）同时捕捉文本的双向上下文关系，克服了传统模型（如LSTM）单向或简单双向拼接的局限性。
- 预训练任务：
- Masked Language Model (MLM)：随机遮盖15%的输入词，模型预测被遮盖的词（学习上下文依赖）。
- Next Sentence Prediction (NSP)：判断两个句子是否是连续的（学习句子间关系）。
- Transformer Encoder架构：
- BERT仅使用Transformer的编码器（Encoder）部分，由多层堆叠的编码器组成，每层包含自注意力机制和前馈神经网络。
二、BERT的层级结构
BERT模型分为输入层、嵌入层、多层编码器。以BERT-Base为例（12层编码器），每一层的作用如下：
1. 输入层（Input Layer）
- 功能：将原始文本转换为模型可处理的输入形式。
- 输入格式：
[CLS]：句首标记，用于分类任务的聚合表示。Token Embeddings：词向量（如WordPiece分词后的词）。Segment Embeddings：区分句子A和句子B（用于NSP任务）。Position Embeddings：位置编码，标记词的位置信息。
2. 嵌入层（Embedding Layer）
- 功能：将输入转换为稠密向量。
- 词嵌入（Token Embeddings）：将词映射到低维向量。
- 位置嵌入（Position Embeddings）：编码词的位置信息。
- 分段嵌入（Segment Embeddings）：区分不同句子（如问答任务中的问题和答案）。
3. 编码器层（Encoder Layers）
每层编码器包含两个核心模块：多头自注意力（Multi-Head Self-Attention）和前馈神经网络（Feed-Forward Network），通过残差连接和层归一化（LayerNorm）优化训练。
-
(1) 多头自注意力机制（Multi-Head Self-Attention）
-
功能：捕捉词与词之间的全局依赖关系。
-
实现：将输入拆分为多个子空间（如12个“头”），每个头独立计算注意力权重，最后拼接结果。
-
公式：
-
Attention(Q,K,V)=softmax(QKT√dk)V
（其中 Q,K,V 是查询、键、值矩阵，( dk ) 是维度）
-
(2) 前馈神经网络（Feed-Forward Network, FFN）
-
功能：对自注意力的输出进行非线性变换。
-
结构：两层全连接层（如中间层维度扩大为4倍），激活函数为GELU/ReLU。
-
(3) 残差连接与层归一化
-
每层输出前应用残差连接 (Output=LayerNorm(x+Sublayer(x)))，缓解梯度消失问题。
4. 各编码器层的特点
- 底层（靠近输入层）：学习基础语法、局部特征（如词性、短语结构）。
- 中层：捕捉句内和句间关系（如指代消解、语义角色）。
- 高层：提取抽象语义（如情感倾向、文本主旨）。
三、BERT的输出
- 最后一层编码器的输出：每个词对应的上下文向量。
[CLS]向量：用于分类任务（如情感分析），聚合全局信息。- 其他词向量：用于序列标注（如命名实体识别）、问答等任务。
环境配置
-
安装 Miniconda：https://www.anaconda.com/docs/getting-started/miniconda/install#power-shell
-
创建虚拟环境：
-
shell conda create -n ai python=3.12 -y # 创建一个虚拟环境 ai conda activate ai # 激活虚拟环境 -
shell # 在官网选择相应Cuda版本，以win11+4060为例 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126 -
验证安装：
-
shell python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())" -
输出：
-
shell 2.6.0+cu126 True # 安装成功
项目结构
my_transformer_demo/
├── config.json
├── data/
│ └── dataset.csv
├── train.py
├── utils.py
配置文件 config.json
{
"model_name": "bert-base-uncased",
"max_length": 128,
"train_batch_size": 8,
"eval_batch_size": 8,
"num_train_epochs": 3,
"learning_rate": 2e-5,
"output_dir": "./models",
"num_labels": 5
}
数据集 data/dataset.csv
sentence,score
"The weather is perfect today!",5
"This restaurant serves awful food.",1
"I'm so happy with my new phone.",5
"The concert was mediocre at best.",3
...
工具函数 utils.py
import pandas as pd
from datasets import Dataset
def load_dataset(path):
df = pd.read_csv(path)
df['label'] = df['score'] - 1 # 评分1~5 → label 0~4
return Dataset.from_pandas(df[['sentence', 'label']])
训练脚本 train.py
import json
import torch
from sklearn.metrics import accuracy_score
from transformers import (
BertTokenizer,
BertForSequenceClassification,
TrainingArguments,
Trainer,
DataCollatorWithPadding,
)
from utils import load_dataset
print(torch.cuda.is_available())
# ========== 设备检测 ==========
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
print(f"🚀 使用 GPU：{torch.cuda.get_device_name(0)}")
print(f"显存占用：{torch.cuda.memory_allocated() / 1024 ** 2:.2f} MB")
else:
print("未检测到 GPU，使用 CPU 训练")
# ========== 加载配置 ==========
with open("config.json") as f:
cfg = json.load(f)
# ========== 加载数据 ==========
tokenizer = BertTokenizer.from_pretrained(cfg["model_name"])
dataset = load_dataset("data/dataset.csv")
def tokenize(example):
return tokenizer(example["sentence"], truncation=True, max_length=cfg["max_length"])
tokenized_dataset = dataset.map(tokenize, batched=True)
tokenized_dataset = tokenized_dataset.train_test_split(test_size=0.2)
# ========== 加载模型 ==========
model = BertForSequenceClassification.from_pretrained(
cfg["model_name"],
num_labels=cfg["num_labels"]
).to(device)
# ========== 评估指标 ==========
def compute_metrics(eval_pred):
logits, labels = eval_pred
preds = logits.argmax(axis=-1)
acc = accuracy_score(labels, preds)
return {"accuracy": acc}
# ========== 训练参数 ==========
args = TrainingArguments(
output_dir=cfg["output_dir"],
per_device_train_batch_size=cfg["train_batch_size"],
per_device_eval_batch_size=cfg["eval_batch_size"],
num_train_epochs=cfg["num_train_epochs"],
learning_rate=cfg["learning_rate"],
save_strategy="epoch",
logging_dir='./logs',
report_to="none",
eval_strategy="epoch",
load_best_model_at_end=True,
)
# ========== 显存监控 Hook ==========
from transformers import TrainerCallback
class PrintMemoryCallback(TrainerCallback):
def on_epoch_begin(self, args, state, control, **kwargs):
if torch.cuda.is_available():
mem = torch.cuda.memory_allocated() / 1024 ** 2
print(f"Epoch {state.epoch:.0f} 开始，当前显存占用：{mem:.2f} MB")
# ========== 启动训练器 ==========
trainer = Trainer(
model=model,
args=args,
train_dataset=tokenized_dataset["train"],
eval_dataset=tokenized_dataset["test"],
tokenizer=tokenizer,
data_collator=DataCollatorWithPadding(tokenizer),
compute_metrics=compute_metrics,
callbacks=[PrintMemoryCallback()]
)
trainer.train()
# ========== 保存模型 ==========
trainer.save_model(cfg["output_dir"])
tokenizer.save_pretrained(cfg["output_dir"])
print("✅ 模型和分词器保存成功！")
使用模型 predict.py
- 模型会保存到
./models/
import sys
import torch
from transformers import BertTokenizer, BertForSequenceClassification
model_path = "models"
tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertForSequenceClassification.from_pretrained(model_path)
def predict(sentence):
inputs = tokenizer(sentence, return_tensors="pt", truncation=True)
with torch.no_grad():
outputs = model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
score = torch.argmax(probs) + 1 # label 0~4 → score 1~5
return score.item()
if __name__ == "__main__":
sentence = " ".join(sys.argv[1:]) or "This is a test."
print(f"Score: {predict(sentence)}")
