BERT-IMDB电影评论情感分类实战：SwanLab可视化训练

基于BERT模型的IMDB电影评论情感分类，是NLP经典的Hello World任务之一。这篇文章我将带大家使用SwanLab、transformers、datasets三个开源工具，完成从数据集准备、代码编写、可视化训练的全过程。观察了一下，中文互联网上似乎很少有能直接跑起来的BERT训练代码和教程，所以也希望这篇文章可以帮到大家。

热爱python的小谢

1197人浏览 · 2025-01-28 08:45:00

热爱python的小谢 · 2025-01-28 08:45:00 发布

基于BERT模型的IMDB电影评论情感分类，是NLP经典的Hello World任务之一。

这篇文章我将带大家使用SwanLab、transformers、datasets三个开源工具，完成从数据集准备、代码编写、可视化训练的全过程。

观察了一下，中文互联网上似乎很少有能直接跑起来的BERT训练代码和教程，所以也希望这篇文章可以帮到大家。

代码：完整代码直接看本文第5节 或 Github
模型与数据集：百度云，提取码: u9gi
实验过程：BERT-SwanLab
SwanLab：swanlab.cn
transformers：github.com/huggingface…
datasets：github.com/huggingface…

1.环境安装

我们需要安装以下这4个Python库：

    transformers>=4.41.0
    datasets>=2.19.1
    swanlab>=0.3.3

一键安装命令：

    pip install transformers datasets swanlab

他们的作用分别是：

transformers：HuggingFace出品的深度学习框架，已经成为了NLP（自然语言处理）领域最流行的训练与推理框架。代码中用transformers主要用于加载模型、训练以及推理。
datasets：同样是HuggingFace出品的数据集工具，可以下载来自huggingface社区上的数据集。代码中用datasets主要用于下载、加载数据集。
swanlab：在线训练可视化和超参数记录工具，官网，可以记录整个实验的超参数、指标、训练环境、Python版本等，并可是化成图表，帮助你分析训练的表现。代码中用swanlab主要用于记录指标和可视化。

本文的代码测试于transformers4.41.0、datasets2.19.1、swanlab==0.3.3，更多库版本可查看SwanLab记录的Python环境。

2.加载BERT模型

BERT模型我们直接下载来自HuggingFace上由Google发布的bert-case-uncased预训练模型。

执行下面的代码，会自动下载模型权重并加载模型：

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments

# 加载预训练的BERT tokenizer
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)

如果国内下载比较慢的话，可以在这个百度云（提取码: u9gi）下载后，把bert-base-uncased文件夹放到根目录，然后改写上面的代码为：

model = AutoModelForSequenceClassification.from_pretrained('./bert-base-uncased', num_labels=2)

3.加载IMDB数据集

IMDB数据集（Internet Movie Database Dataset）是自然语言处理（NLP）领域中一个非常著名和广泛使用的数据集，主要应用于文本情感分析任务。

IMDB数据集源自全球最大的电影数据库网站Internet Movie Database（IMDb），该网站包含了大量的电影、电视节目、纪录片等影视作品信息，以及用户对这些作品的评论和评分。数据集包括50,000条英文电影评论，这些评论被标记为正面或负面情感，用以进行二分类任务。其中，25,000条评论被分配为训练集，另外25,000条则作为测试集。训练集和测试集都保持了平衡的正负样本比例，即各含50%的正面评论和50%的负面评论.

我们同样直接下载HuggingFace上的imdb数据集，执行下面的代码，会自动下载数据集并加载：

from datasets import load_dataset

# 加载IMDB数据集
dataset = load_dataset('imdb')

如果国内下载比较慢的话，可以在这个百度云（提取码: u9gi）下载后，把imdb文件夹放到根目录，然后改写上面的代码为：

dataset = load_dataset('./imdb')

4.集成SwanLab

因为swanlab已经和transformers框架做了集成，所以将SwanLabCallback类传入到trainer的callbacks参数中即可实现实验跟踪和可视化：

from swanlab.integration.huggingface import SwanLabCallback

# 设置swanlab回调函数
swanlab_callback = SwanLabCallback()

...

# 定义Transformers Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['test'],
    # 传入swanlab回调函数
    callbacks=[swanlab_callback],
)

想了解更多关于SwanLab的知识，请看SwanLab官方文档。

5.开始训练！

训练过程看这里：BERT-SwanLab。

在首次使用SwanLab时，需要去官网注册一下账号，然后在用户设置复制一下你的API Key。

然后在终端输入swanlab login:

swanlab login

把API Key粘贴进去即可完成登录，之后就不需要再次登录了。

完整的训练代码：

"""
用预训练的Bert模型微调IMDB数据集，并使用SwanLabCallback回调函数将结果上传到SwanLab。
IMDB数据集的1是positive，0是negative。
"""

import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from swanlab.integration.huggingface import SwanLabCallback
import swanlab

def predict(text, model, tokenizer, CLASS_NAME):
    inputs = tokenizer(text, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class = torch.argmax(logits).item()

    print(f"Input Text: {text}")
    print(f"Predicted class: {int(predicted_class)} {CLASS_NAME[int(predicted_class)]}")
    return int(predicted_class)

# 加载IMDB数据集
dataset = load_dataset('imdb')

# 加载预训练的BERT tokenizer
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')

# 定义tokenize函数
def tokenize(batch):
    return tokenizer(batch['text'], padding=True, truncation=True)

# 对数据集进行tokenization
tokenized_datasets = dataset.map(tokenize, batched=True)

# 设置模型输入格式
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
tokenized_datasets.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])

# 加载预训练的BERT模型
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)

# 设置训练参数
training_args = TrainingArguments(
    output_dir='./results',
    eval_strategy='epoch',
    save_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    logging_first_step=100,
    # 总的训练轮数
    num_train_epochs=3,
    weight_decay=0.01,
    report_to="none",
    # 单卡训练
)

CLASS_NAME = {0: "negative", 1: "positive"}

# 设置swanlab回调函数
swanlab_callback = SwanLabCallback(project='BERT',
                                   experiment_name='BERT-IMDB',
                                   config={'dataset': 'IMDB', "CLASS_NAME": CLASS_NAME})

# 定义Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['test'],
    callbacks=[swanlab_callback],
)

# 训练模型
trainer.train()

# 保存模型
model.save_pretrained('./sentiment_model')
tokenizer.save_pretrained('./sentiment_model')

# 测试模型
test_reviews = [
    "I absolutely loved this movie! The storyline was captivating and the acting was top-notch. A must-watch for everyone.",
    "This movie was a complete waste of time. The plot was predictable and the characters were poorly developed.",
    "An excellent film with a heartwarming story. The performances were outstanding, especially the lead actor.",
    "I found the movie to be quite boring. It dragged on and didn't really go anywhere. Not recommended.",
    "A masterpiece! The director did an amazing job bringing this story to life. The visuals were stunning.",
    "Terrible movie. The script was awful and the acting was even worse. I can't believe I sat through the whole thing.",
    "A delightful film with a perfect mix of humor and drama. The cast was great and the dialogue was witty.",
    "I was very disappointed with this movie. It had so much potential, but it just fell flat. The ending was particularly bad.",
    "One of the best movies I've seen this year. The story was original and the performances were incredibly moving.",
    "I didn't enjoy this movie at all. It was confusing and the pacing was off. Definitely not worth watching."
]

model.to('cpu')
text_list = []
for review in test_reviews:
    label = predict(review, model, tokenizer, CLASS_NAME)
    text_list.append(swanlab.Text(review, caption=f"{label}-{CLASS_NAME[label]}"))

![bert-swanlab-2.png](https://p9-juejin.byteimg.com/tos-cn-i-k3u1fbpfcp/14c57e319ec441339d5b0e2c5b8bba6e~tplv-k3u1fbpfcp-jj-mark:0:0:0:0:q75.image#?w=2724&h=970&s=181947&e=png&b=fefefe)
if text_list:
    swanlab.log({"predict": text_list})

swanlab.finish()

训练可视化过程：

训练大概需要6G左右的显存，我在一块3090上跑了，1个epoch大概要12～13分钟时间。

训练的推理结果：

这里我生成了10个比较简单的测试文本，微调后的BERT模型基本都能答对。

至此，我们顺利完成了用BERT预训练模型微调IMDB数据的训练过程～

如何学习大模型

下面这些都是我当初辛苦整理和花钱购买的资料，现在我已将重要的AI大模型资料包括市面上AI大模型各大白皮书、AGI大模型系统学习路线、AI大模型视频教程、实战学习，等录播视频免费分享出来，需要的小伙伴可以扫取。

现在社会上大模型越来越普及了，已经有很多人都想往这里面扎，但是却找不到适合的方法去学习。

作为一名资深码农，初入大模型时也吃了很多亏，踩了无数坑。现在我想把我的经验和知识分享给你们，帮助你们学习AI大模型，能够解决你们学习中的困难。

👉学会后的收获：👈
• 基于大模型全栈工程实现（前端、后端、产品经理、设计、数据分析等），通过这门课可获得不同能力；

• 能够利用大模型解决相关实际项目需求：大数据时代，越来越多的企业和机构需要处理海量数据，利用大模型技术可以更好地处理这些数据，提高数据分析和决策的准确性。因此，掌握大模型应用开发技能，可以让程序员更好地应对实际项目需求；

• 基于大模型和企业数据AI应用开发，实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能，学会Fine-tuning垂直训练大模型（数据准备、数据蒸馏、大模型部署）一站式掌握；

• 能够完成时下热门大模型垂直领域模型训练能力，提高程序员的编码能力：大模型应用开发需要掌握机器学习算法、深度学习框架等技术，这些技术的掌握可以提高程序员的编码能力和分析能力，让程序员更加熟练地编写高质量的代码。

👉 福利来袭CSDN大礼包：《2025最全AI大模型学习资源包》免费分享，安全可点 👈

一、AGI大模型系统学习路线

很多人学习大模型的时候没有方向，东学一点西学一点，像只无头苍蝇乱撞，我下面分享的这个学习路线希望能够帮助到你们学习AI大模型。

在这里插入图片描述

二、AI大模型视频教程

在这里插入图片描述

三、AI大模型各大学习书籍!

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

四、AI大模型各大场景实战案例

在这里插入图片描述

五、AI大模型面试题库
在这里插入图片描述

五、结束语

学习AI大模型是当前科技发展的趋势，它不仅能够为我们提供更多的机会和挑战，还能够让我们更好地理解和应用人工智能技术。通过学习AI大模型，我们可以深入了解深度学习、神经网络等核心概念，并将其应用于自然语言处理、计算机视觉、语音识别等领域。同时，掌握AI大模型还能够为我们的职业发展增添竞争力，成为未来技术领域的领导者。

再者，学习AI大模型也能为我们自己创造更多的价值，提供更多的岗位以及副业创收，让自己的生活更上一层楼。

因此，学习AI大模型是一项有前景且值得投入的时间和精力的重要选择。