简介:基于BERT的情感分析项目源码,面向NLP初学者与中级学习者,契合IMDB影评正面/负面二分类任务,适用于课程设计、毕业设计或Transformer实践入门。压缩包体积仅4KB,包含5个文件:4个Python脚本与1个说明文本,代码量精简但覆盖了模型构建、数据预处理、训练、测试与环境适配等关键环节。其中code.py为主程序,Examination.py、GPU TEST.py、PyTorch TEST.py分别用于结果检验、GPU模式运行及PyTorch框架下的测试,方便读者在不同计算资源下切换验证,也便于对比不同运行方式的输出差异。项目源码已本地编译通过,难度适中,并经过助教审定,运行稳定性有一定保障。配套使用说明txt可帮助理清各文件用途及执行顺序,降低配置成本,同时为后续二次开发提供简单指引。已有226人学习,适合希望快速掌握BERT微调逻辑、情感分类流程以及Transformer模型从加载到评估的全链路参考的读者,可作为理解预训练模型文本分类应用的直观案例。
1. 用 BERT 给 IMDB 影评做情感分类,这个 Python 项目到底值不值得跑?
如果你手里正好有一份"基于 BERT 模型的情感分析项目"的 Python 源码,目标是给 IMDB 影评分出正面或负面,那你大概率是在做 NLP 入门到进阶的必修课。这个项目听起来像教科书示例,但它实际上覆盖了从文本预处理、预训练模型加载、微调训练到模型评估的完整链路,是理解 Transformers 架构落地最直接的方式之一。IMDB 数据集有 5 万条带标签影评,正负各半,规模适中,既不会让你的显卡一夜烧穿,也足够让你感受到过拟合、显存溢出和训练时间失控这些真实工程问题。
先说结论:这个项目非常适合作为你第一个"真正运行起来的深度学习 NLP 任务"。它不需要你从零实现 Attention,只需要你理解 BERT 怎么用、数据怎么喂、参数怎么调。我会按照数据准备、模型加载、训练、评估、避坑、部署的顺序,把完整可复现的代码和参数一起给你,新手照着跑能出结果,熟手也能从参数边界里找到调优空间。
2. 数据准备与预处理:IMDB 原始文本是怎么变成 BERT 能吃的张量的
2.1 IMDB 数据集的获取与目录结构
IMDB 影评数据集最常见的获取方式是直接使用 TensorFlow Datasets 或者 Hugging Face Datasets 库。但如果你拿到的 Python 源码里自带数据文件夹,通常会是aclImdb/train/pos、aclImdb/train/neg、aclImdb/test/pos、aclImdb/test/neg这样的目录结构。每个目录下是一堆 .txt 文件,每个文件里就是一条影评原文。
我一般建议优先用 Hugging Face 的datasets库加载,因为它会自动帮你划分训练集和测试集,还自带unsup无监督数据。手动读目录的方式很容易在文件编码和路径拼接上出问题,Windows 和 Linux 的路径分隔符还不一样。以下代码两种方式都兼容,你根据手头源码的实际情况选一种:
# 方式一:使用 Hugging Face datasets 库(推荐) from datasets import load_dataset dataset = load_dataset("imdb") print(dataset) # DatasetDict({ # train: Dataset({features: ['text', 'label'], num_rows: 25000}) # test: Dataset({features: ['text', 'label'], num_rows: 25000}) # unsupervised: Dataset({features: ['text', 'label'], num_rows: 50000}) # }) # 方式二:手动读取原始文件目录 import os from glob import glob def read_imdb_folder(pos_dir, neg_dir): texts, labels = [], [] for file_path in glob(os.path.join(pos_dir, "*.txt")): with open(file_path, "r", encoding="utf-8") as f: texts.append(f.read()) labels.append(1) # 正面为 1 for file_path in glob(os.path.join(neg_dir, "*.txt")): with open(file_path, "r", encoding="utf-8") as f: texts.append(f.read()) labels.append(0) # 负面为 0 return texts, labels这段代码的逻辑很简单:Hugging Face 方式直接拿到带 label 的数据集,手动方式则把正负样本分别读入并标记为 1 和 0。注意 IMDB 原始数据里有些影评文件是空的或者极短(只有几个字符),加载后最好过滤掉长度小于 20 个字符的样本,否则 tokenizer 之后可能只剩 [CLS] 和 [SEP],模型无从判断情感倾向。
2.2 BERT Tokenizer 的分词原理与加载方式
BERT 用的是 WordPiece 分词,它会把 "unhappy" 拆成 "un" + "happy",把 "playing" 拆成 "play" + "##ing"。这意味着你不能直接用 Python 的split()按空格切词,必须用配套的 tokenizer。加载方式如下:
from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") # 或者使用源码里可能自带的本地 tokenizer 目录 # tokenizer = AutoTokenizer.from_pretrained("./bert-base-uncased/") text = "This movie is absolutely fantastic!" tokens = tokenizer.tokenize(text) ids = tokenizer.encode(text, add_special_tokens=True) print(tokens) # ['this', 'movie', 'is', 'absolutely', 'fan', '##tastic', '!'] print(ids) # [101, 2023, 3185, 2003, 7078, 6353, 11432, 999, 102]101是 [CLS] 的 ID,102是 [SEP] 的 ID。每个样本都会在开头加 [CLS]、结尾加 [SEP],模型最后只取 [CLS] 位置的输出做分类。这里有几个隐藏细节:小写 uncased 模型会自动把所有字母转成小写,如果你做的是英文影评分类,uncased 足够;但如果影评里大量出现专有名词(如电影角色名 "Kane"),cased 模型可能略好一点,这需要你实验对比。
2.3 序列截断与填充:max_length 选 128 还是 512
IMDB 影评长度差异很大,有的只有一句话,有的能写几千词。BERT 的输入长度上限是 512 个 token,但直接把所有样本都填充到 512 会严重浪费算力,因为大部分影评集中在 100-300 token 之间。这里我给出一个统计方法,你先看分布再定参数:
import numpy as np from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") lengths = [] # 假设 texts 是你的原始影评列表 for text in texts[:5000]: length = len(tokenizer.encode(text, add_special_tokens=True)) lengths.append(length) lengths = np.array(lengths) print(f"平均长度: {lengths.mean():.1f}") print(f"中位数: {np.median(lengths):.1f}") print(f"90分位: {np.percentile(lengths, 90):.1f}") print(f"95分位: {np.percentile(lengths, 95):.1f}")大多数情况下 90 分位在 250-350 之间。我一般建议取 256 作为 max_length,如果显存紧张也可以取 128,代价是丢失长影评的尾部信息。截断策略用truncation=True,BERT 默认截断尾部,因为头部包含 [CLS] 和开头几句话,通常信息密度更高,这也符合阅读习惯——很少有人看了开头三句话还判断不出电影好坏趋势。
2.4 PyTorch Dataset 封装与 DataLoader 的参数选择
Tokenizer 预处理后,我们需要把数据封装成 PyTorch 的 Dataset,再交给 DataLoader 分批训练。这一步的细节决定了训练时数据喂给模型的效率:
import torch from torch.utils.data import Dataset, DataLoader class IMDBDataset(Dataset): def __init__(self, texts, labels, tokenizer, max_length=256): self.texts = texts self.labels = labels self.tokenizer = tokenizer self.max_length = max_length def __len__(self): return len(self.texts) def __getitem__(self, idx): text = self.texts[idx] label = self.labels[idx] encoding = self.tokenizer( text, truncation=True, padding="max_length", max_length=self.max_length, return_tensors="pt" ) return { "input_ids": encoding["input_ids"].squeeze(0), "attention_mask": encoding["attention_mask"].squeeze(0), "label": torch.tensor(label, dtype=torch.long) } # 实例化数据集 train_dataset = IMDBDataset(train_texts, train_labels, tokenizer, max_length=256) val_dataset = IMDBDataset(val_texts, val_labels, tokenizer, max_length=256) # DataLoader 参数说明 train_loader = DataLoader( train_dataset, batch_size=16, # 显存不够就降到 8,但学习率也要相应调整 shuffle=True, num_workers=4, # Windows 上建议设为 0,否则会报多进程错误 pin_memory=True # GPU 训练时加速 host 到 device 的传输 ) val_loader = DataLoader( val_dataset, batch_size=32, # 验证阶段可以不反向传播,batch 可以大一点 shuffle=False, num_workers=4, pin_memory=True )在这里我踩过好几个坑,先说两个最典型的。第一,padding="max_length"会把所有样本都填到 256,这比动态 padding(padding=True)更浪费计算,但好处是维度固定、不易出错,新手阶段先用 max_length 更稳妥;等跑通之后再改成动态 padding 提升训练速度。第二,attention_mask一定要传,它是告诉 BERT"哪些位置是真实 token、哪些是 padding",不传的话模型会把 padding 也当文本处理,训练出来准确率会掉 3-5 个百分点。"padding 也会被模型当真" 这个坑在源码里最隐蔽,很多人模型代码改来改去没发现问题,最后才发现是 mask 没传。
3. 模型加载与训练流程:从bert-base-uncased到分类头的微调
3.1 为什么用AutoModelForSequenceClassification而不是裸 BERT
拿到源码后你可能会看到两种写法。一种是从transformers直接加载AutoModelForSequenceClassification,另一种是先加载AutoModel再在 [CLS] 输出后面自己接一个全连接层。前者是官方推荐做法,后者更灵活但容易写错。
AutoModelForSequenceClassification其实就是在 BERT 的 [CLS] 输出上接了一个线性分类层。原始 BERT 输出是 768 维向量(base 版本),分类层把它映射到num_labels个类别上。对 IMDB 情感分类来说就是 2 类,所以分类层的权重矩阵是[768, 2]。直接使用这个封装类的好处是它自动处理了 dropout、pooler 输出和 loss 计算,你不需要手动写nn.CrossEntropyLoss的逻辑,前向传播时传入labels就会自动返回 loss。
还有一种情况,你的源码里可能加载的是bert-base-uncased的 PyTorch 版本本地目录。这时候加载路径就是本地文件夹路径而不是模型名,这一点在from_pretrained里会自动识别——你传一个存在的目录路径它就按本地加载处理。
3.2 训练超参数:batch_size、学习率、epoch 的合理区间
BERT 微调不是从零训练,我们不能用大学习率。预训练模型已经学到了通用语言表示,微调只需要在它的基础上"稍微调整"。常用配置区间如下:
| 参数 | 推荐值 | 说明 |
|---|---|---|
| batch_size | 8-16 | 太大容易 OOM,太小收敛慢 |
| learning_rate | 2e-5 到 5e-5 | 超过 5e-5 容易灾难性遗忘 |
| epoch | 2-4 | IMDB 上 3 epoch 基本到顶,再多过拟合 |
| warmup_steps | 训练步数的 10% | 避免初期 loss 震荡 |
| weight_decay | 0.01 | 对分类头有效,对 BERT 主体影响小 |
| max_grad_norm | 1.0 | 梯度裁剪,防止 loss 突变 |
3.3 完整训练脚本:训练循环、梯度累积与模型保存
下面给出一个完整的训练脚本,注释里写明了每一步在做什么。这个脚本可以直接跑,但你要注意把优化器和学习率调度器的参数与你的显存和数据集规模匹配:
import torch from transformers import AutoModelForSequenceClassification, AdamW, get_linear_schedule_with_warmup from tqdm import tqdm import os model = AutoModelForSequenceClassification.from_pretrained( "bert-base-uncased", num_labels=2 ) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) total_steps = len(train_loader) * 3 # 3 epochs optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-8) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=int(total_steps * 0.1), num_training_steps=total_steps ) criterion = torch.nn.CrossEntropyLoss() # 梯度累积步数:batch_size=8 想达到 16 的效果就在每 2 个 batch 后更新一次 accumulation_steps = 1 model.zero_grad() for epoch in range(3): model.train() total_loss = 0 for step, batch in enumerate(tqdm(train_loader)): input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["label"].to(device) outputs = model( input_ids=input_ids, attention_mask=attention_mask, labels=labels ) loss = outputs.loss total_loss += loss.item() # 梯度累积:小显存模拟大 batch,每 accumulation_steps 步更新一次权重 loss = loss / accumulation_steps loss.backward() if (step + 1) % accumulation_steps == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() scheduler.step() model.zero_grad() avg_loss = total_loss / len(train_loader) print(f"Epoch {epoch + 1} 平均训练损失: {avg_loss:.4f}") # 每个 epoch 结束后保存一次,方便回滚到最佳状态 checkpoints_dir = "./checkpoints" os.makedirs(checkpoints_dir, exist_ok=True) model.save_pretrained(f"{checkpoints_dir}/bert-imdb-epoch-{epoch + 1}") tokenizer.save_pretrained(f"{checkpoints_dir}/bert-imdb-epoch-{epoch + 1}") print("训练完成!")记忆一个调整技巧:如果你发现训练损失下降很慢,先看是不是学习率太小(低于 1e-5);如果训练损失下降但验证损失上升,说明已经过拟合,应该马上停止而不是等 3 个 epoch 跑完。你可以在每个 epoch 结束后跑一次验证集,把验证准确率最高的那个 epoch 的模型作为最终产物。
3.4 为什么 3 个 epoch 就够了:BERT 微调的收敛规律
很多第一次跑 BERT 的人都有疑虑:3 个 epoch 是不是太少了?CNN 训练 MNIST 都要几十个 epoch。这是因为 BERT 的参数已经在海量语料上预训练过,它的词向量和注意力模式已经非常成熟。微调阶段要做的是"适配"而不是"学习"。
IMDB 这个任务本质上就是判断文本的情感极性,BERT 在预训练阶段已经见过大量类似表达,微调只需要让分类头学会"把 [CLS] 输出映射到正负标签"。从实验曲线上看,第一个 epoch 验证准确率通常就能到 85% 以上,第二个 epoch 到 90% 以上,第三个 epoch 可能只提升 0.5-1 个百分点,第四个 epoch 开始验证指标停滞或下降。这个时候与其继续训练,不如去做模型融合或者尝试不同的 max_length。实际训练中我们还发现,使用distilbert-base-uncased这个轻量版本,训练速度快一倍,准确率只下降 1-2 个百分点,对于新手验证流程来说是性价比更高的选择。
4. 评估指标与模型验证:准确率不是唯一标准,F1 和混淆矩阵才是真相
4.1 分类报告:精度、召回率、F1 怎么解读
IMDB 的正负样本是完全均衡的(各 25,000 条),所以准确率看起来是个合理指标。但实际预测时模型可能会偏向某一类,尤其是当你用bert-base-uncased在 IMDB 训练 3 个 epoch 后,模型往往对负面评论的召回率更高。为什么?因为负面评论里往往有更强烈的情绪词("terrible" "waste" "boring"),这些词的特征更突出,模型更容易捕捉。
我建议你至少打印一份完整的分类报告:
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score import torch def evaluate(model, val_loader, tokenizer, max_length=256): model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for batch in val_loader: input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["label"].to(device) outputs = model(input_ids=input_ids, attention_mask=attention_mask) logits = outputs.logits preds = torch.argmax(logits, dim=-1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) print("准确率:", accuracy_score(all_labels, all_preds)) print("\n分类报告:\n", classification_report(all_labels, all_preds, target_names=["负面", "正面"])) # 混淆矩阵可以辅助判断模型偏见 print("混淆矩阵:\n", confusion_matrix(all_labels, all_preds)) # 假设你加载的是第 2 个 epoch 的 checkpoint model = AutoModelForSequenceClassification.from_pretrained("./checkpoints/bert-imdb-epoch-2") model.to(device) evaluate(model, val_loader, tokenizer)分类报告里重点关注"正面"类别的召回率——如果模型把 20% 的正面评论误判为负面,而负面评论召回率却很高,说明模型对"中性偏正面"的温和好评识别力不足。IMDB 里大量 5-7 分的影评措辞温和,比如 "It was decent but not great",这类样本正是模型的盲区。
4.2 过拟合信号:训练损失和验证指标的分歧
你可以在训练过程中记录每个 epoch 的训练损失和验证准确率,画成曲线看走势。正常情况是训练损失下降的同时验证准确率上升;一旦训练损失还在降但验证准确率停滞或下降,过拟合就开始了。
对 BERT 微调来说,过拟合的表现不像 CNN 那么剧烈。你可能看到训练损失降到 0.1 以下,但验证准确率停留在 91% 上不去。这时除了增加 dropout 外,还有一个有意思的做法:冻结 BERT 前几层的参数,只训练后面的层和分类头。实践证明这能显著提升验证集表现,也是解决小数据集过拟合的常见方案:
# 冻结前 8 层 Transformer 层,只训练后 4 层和分类头 for name, param in model.named_parameters(): if "bert.encoder.layer" in name: layer_num = int(name.split("layer.")[1].split(".")[0]) if layer_num < 8: param.requires_grad = False4.3 样例可视化:哪些评论被分错了
光看指标不够,你还需要知道模型具体在哪些样例上出错。跑一遍预测,把预测错误的样本打印出来,大概率会发现以下规律:模型把含有 "not bad" "hard to watch but worth it" 这类转折句式的评论判错,把反讽、隐喻和大量电影梗的评论判错。这属于 BERT 的已知盲区,它本质上是概率模型,不理解"反讽"这种需要外部知识才能解码的修辞手段。
def show_misclassified(model, val_dataset, tokenizer, num_samples=10): model.eval() device = next(model.parameters()).device shown = 0 for text, true_label in zip(val_dataset.texts, val_dataset.labels): encoding = tokenizer(text, truncation=True, padding="max_length", max_length=256, return_tensors="pt") input_ids = encoding["input_ids"].to(device) attention_mask = encoding["attention_mask"].to(device) with torch.no_grad(): logits = model(input_ids=input_ids, attention_mask=attention_mask).logits pred_label = torch.argmax(logits, dim=-1).item() if pred_label != true_label: print(f"原文: {text[:200]}...") print(f"真实: {'正面' if true_label == 1 else '负面'} | 预测: {'正面' if pred_label == 1 else '负面'}") print("---") shown += 1 if shown >= num_samples: break这个函数不单是演示用,在实际项目中你可以把它写成一个脚本,每次模型更新后跑一遍最近一批真实评论,人工检查模型的错误模式有没有变化。比如模型在某个时间段突然大量误判包含 "director" 的评论,很可能是训练数据里这类样本太少或者有标注错误。
5. 避坑指南:BERT 情感分析项目里最常见的 5 个野路子问题
5.1 显存不足(OOM)报错:RuntimeError: CUDA out of memory
现象:训练到一半直接报CUDA out of memory,有时是第一个 epoch 刚开始就崩,有时是跑了几百个 step 才崩。
原因:IMDB 虽然单条文本经过max_length=256截断后不算大,但如果你设了batch_size=32又不小心开启了gradient_checkpointing的负优化,或者同一时间在 GPU 上跑了其他程序(比如 Jupyter notebook 里上一个 cell 的模型还没释放),显存就会不够。另外很多人忽略了一个细节:from_pretrained时会加载完整 BERT 权重,如果你的显存只有 4GB,batch_size=16+max_length=512几乎必崩,因为 backpropagation 需要保存前向传播中的中间激活值,这部分显存占用远超模型参数本身。
解决:先把batch_size降到 8 或 4;如果还崩,把max_length从 512 降到 256。记住梯度累积可以等效增大 batch 而不会增加显存——这就是我在前面代码里写accumulation_steps的原因。还可以在训练前执行torch.cuda.empty_cache()清掉缓存碎片。如果显存实在放不下,最后的手段是启用model.gradient_checkpointing_enabled(),它会用"前向时丢弃中间激活、反向时重新计算"来换取显存,速度变慢但能跑。
5.2 训练是能跑,但 loss 一直不降,准确率卡在 50%
现象:训练了 2 个 epoch,loss 在 0.7 附近来回晃动,验证准确率 51% 左右,跟瞎猜差不多。
原因:标签和输入不匹配是最常见的原因——return_tensors="pt"返回的 input_ids 是多一维的形状([1,seq_len]),直接传给模型时模型把它当成了 batch_size=1 的一个样本组,训练逻辑变成了逐样本梯度更新,且 attention_mask 与 input_ids 在 batch 维度上没对齐。另一个可能原因是标签把 1 和 0 弄反了,训练集里 pos 目录下的文本被标成了 0,模型在反向学习。"loss 毫无反应先别调学习率,先检查数据流" 是我反复给自己强调的。解决:打印一个 batch 出来看,input_ids.shape应该是[16, 256]而不是[16, 1, 256];在__getitem__里用.squeeze(0)去掉多余的 batch 维度。然后手动抽查 5 条训练样本,打印 text 和 label 的对应关系,确认 1 对应正面评价("amazing" "love" 等词),0 对应负面("terrible" "hate" 等词)。
5.3 验证集准确率正常,但拿真实评论一预测全模糊
现象:训练完测试集准确率有 91%,但你拿着一条现实中的影评去测,模型输出的概率接近 50%,或者对明显很正面的评论给出 55% 的正面概率。
原因:IMDB 数据集里的影评都是用户写的长评,结构相对完整。而现实场景中你可能会输入短评、一句话吐槽,甚至带 emoji 和 @ 符号的文本。这类文本的分布与训练集差异很大。另一个原因是 tokenizer 会删除或截断特殊字符,BERT 对 emoji 的处理是把它们变成[UNK],等于丢掉了信息。解决:做一个针对短文本的二次微调。收集 2000-5000 条短影评(包括社交平台短评),用训练好的模型继续微调 1 个 epoch,学习率降到 1e-5。这招在实战中很管用,它让模型从"长文阅读模式"切换到"短文本快速判断模式"。另外在预测时不要设max_length=256,用truncation=True+ 文本原长输入,避免短文本被 padding 成一大堆没意义的 [PAD] token。
5.4 效果还行,但训练和推理速度异常慢
现象:3 个 epoch 要跑 6 个小时,3000 条评论的预测要等 10 分钟。
原因:最常见的是没有用 GPU 而是在 CPU 上训练。然后检查num_workers是否在 Windows 上设置为大于 0——这会导致每个 epoch 启动时反复加载数据,甚至会报BrokenPipeError。还有一个隐性原因:把验证集的 shuffle 设为True,每次验证都打乱顺序,导致每次都重新分配内存,白白消耗时间。
解决:确认torch.cuda.is_available()为 True 且 model 已经.to(device);Windows 上num_workers设 0。训练时在你的 transformer 配置里添加use_cache=False,对 decoder-only 模型收益更大。推理时开启torch.no_grad()并且用torch.inference_mode()替代前者的平均加速比更高。另外如果你用的是bert-base-uncased,可以考虑换distilbert-base-uncased做推理——权重小一半,速度提升近 2 倍,分类效果下降不到 2 个百分点,在真实部署场景里这是更划算的选择。
5.5 部署加载模型时from_pretrained报错无法找到模型
现象:在本地测试完,把 checkpoint 目录打包发给同事,同事加载时报OSError: Unable to load weights from pytorch_model.bin,或者提示Some weights of the model checkpoint were not used。
原因:save_pretrained是保存到了本地目录,但from_pretrained加载时如果路径写成了裸文件名(比如"bert-imdb-epoch-2"而不是"./checkpoints/bert-imdb-epoch-2"),它会被误认为要去 Hugging Face Hub 下载同名模型。更坑的是Some weights were not used警告——这不一定报错,但如果分类头维度不对配(你保存的是 2 分类模型,加载时又传了num_labels=5),就会初始化失败而不自知。
解决:永远用绝对路径或带./的相对路径加载本地模型目录,确保目录下有config.json、pytorch_model.bin(或model.safetensors)、vocab.txt和tokenizer_config.json四个文件。加载时不再传num_labels,直接用AutoModelForSequenceClassification.from_pretrained("./checkpoints/bert-imdb-epoch-2"),让模型从 config 里读标签数。顺带提一句,新版 transformers 默认保存 safetensors 格式,加载时要求 transformers 版本不低于 4.30,不然会报二进制解析错误,这个坑遇到过的人深有体会。
6. 进阶用法与部署验证:把训练好的模型变成能用的推理服务
训练和评估跑通之后,你手里就有了一个能用的bert-imdb模型文件目录。现在的问题是怎么让"别人也能用"——不是指发源码,而是指可以拿一条新影评进去,得到有业务意义的情感判断。下面给出一个完整的推理脚本,比你在源码里可能见到的裸预测多做了三件事:软标签输出、置信度提示、以及一个简单的批处理模式。
import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer class IMDBClassifier: def __init__(self, model_path): self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.tokenizer = AutoTokenizer.from_pretrained(model_path) self.model = AutoModelForSequenceClassification.from_pretrained(model_path) self.model.to(self.device) self.model.eval() def predict_prob(self, text: str): encoding = self.tokenizer( text, truncation=True, max_length=256, padding="max_length", return_tensors="pt" ) input_ids = encoding["input_ids"].to(self.device) attention_mask = encoding["attention_mask"].to(self.device) with torch.inference_mode(): logits = self.model(input_ids=input_ids, attention_mask=attention_mask).logits probs = torch.softmax(logits, dim=-1).cpu().numpy()[0] return probs def predict_label(self, text: str): probs = self.predict_prob(text) label = int(probs.argmax()) confidence = probs[label] sentiment = "正面" if label == 1 else "负面" print(f"情感倾向: {sentiment}(置信度: {confidence:.4f})") return label, float(confidence) def predict_batch(self, texts: list): results = [] for text in texts: label, confidence = self.predict_label(text) results.append({"text": text, "label": label, "confidence": confidence}) return results # 使用示例 classifier = IMDBClassifier("./checkpoints/bert-imdb-epoch-3") sample_text = "This movie was a masterpiece. The acting was phenomenal and the plot kept me on the edge of my seat!" print(classifier.predict_prob(sample_text)) # [0.0012, 0.9988] -> 98.88% 正面最后说几个进阶方向,供你跑通这个项目之后参考:第一,尝试在验证集上做阈值调优,把默认的 0.5 概率阈值改为 0.6 或 0.4,观察 precision-recall 的变化,这在业务对其中一类要求更高时很有用。第二,用DistilBERT或ALBERT替换bert-base-uncased做模型蒸馏对比,你会看到参数量下降后的精度损失,也能体会 "训练速度提升 2 倍、显存占用减半" 是什么体验。第三,把模型导出为 ONNX 格式,用onnxruntime做推理,单条样本的延迟可以从 50ms 降到 10ms 级别——如果你的影评分析要走向实时 API 服务,这一步是必经之路。
说回我自己的经验:每次把这个流程从头到尾跑通一遍,我都会发现新的坏习惯——有时是忘了设置dropout的training状态,有时是推理时忘了model.eval()导致结果随机浮动。BERT 情感分析这个项目真正教会我的不是from_pretrained怎么用,而是"数据流检查的顺序":先确认数据形状,再确认标签方向,然后看 loss 曲线,最后才动模型结构。按照这套顺序排查,90% 的问题都能在半分钟内定位。希望这篇笔记能帮你把这个项目跑出 91% 以上的测试准确率,并且让你之后在面对其他文本分类任务时,知道从哪里下手查问题。
本文还有配套的精品资源,点击获取