Jupyter Notebook超60万条数据集情感分析提速及中断问题求助

核心优化代码 #

基础加速版本（适合无中断风险场景） #

import pandas as pd
from camel_tools.sentiment import SentimentAnalyzer
import torch

# 自动检测GPU并加载模型到GPU（核心提速点）
device = "cuda" if torch.cuda.is_available() else "cpu"
sa = SentimentAnalyzer("CAMeL-Lab/bert-base-arabic-camelbert-da-sentiment", device=device)

# 批量大小根据GPU显存调整：显存充足设256/512，显存有限设32/64
BATCH_SIZE = 128
full_text = dataset['clean_text'].tolist()
total_samples = len(full_text)
sentiments = []

# 批量处理样本
for i in range(0, total_samples, BATCH_SIZE):
    batch_text = full_text[i:i+BATCH_SIZE]
    # 批量预测（比逐条推理效率高10-100倍）
    batch_sentiments = sa.predict(batch_text)
    sentiments.extend(batch_sentiments)
    
    # 每处理100个批次打印进度，避免逐行拖慢速度
    if (i // BATCH_SIZE) % 100 == 0:
        print(f"已完成 {i + len(batch_text)}/{total_samples} 条样本")

# 将结果写入原数据集
dataset['sentiments'] = sentiments

断点续传版本（适合易中断的长时任务） #

import pandas as pd
from camel_tools.sentiment import SentimentAnalyzer
import torch
import os

# 断点文件路径，用于保存已处理进度
CHECKPOINT_PATH = "sentiment_progress.pkl"
device = "cuda" if torch.cuda.is_available() else "cpu"
sa = SentimentAnalyzer("CAMeL-Lab/bert-base-arabic-camelbert-da-sentiment", device=device)

BATCH_SIZE = 128
full_text = dataset['clean_text'].tolist()
total_samples = len(full_text)
sentiments = []
processed_num = 0

# 从断点恢复进度
if os.path.exists(CHECKPOINT_PATH):
    checkpoint = pd.read_pickle(CHECKPOINT_PATH)
    sentiments = checkpoint['sentiments']
    processed_num = checkpoint['processed_num']
    print(f"从断点恢复，已处理 {processed_num} 条样本")

# 继续处理剩余样本
for i in range(processed_num, total_samples, BATCH_SIZE):
    batch_text = full_text[i:i+BATCH_SIZE]
    batch_sentiments = sa.predict(batch_text)
    sentiments.extend(batch_sentiments)
    processed_num = i + len(batch_text)
    
    # 每100批次保存一次断点
    if (i // BATCH_SIZE) % 100 == 0:
        print(f"已完成 {processed_num}/{total_samples} 条样本")
        pd.to_pickle({'sentiments': sentiments, 'processed_num': processed_num}, CHECKPOINT_PATH)

# 最终写入结果并清理断点文件
dataset['sentiments'] = sentiments
if os.path.exists(CHECKPOINT_PATH):
    os.remove(CHECKPOINT_PATH)

关键优化说明 #

• GPU并行加速：显式将模型加载到GPU，比CPU推理速度提升10-100倍，是核心提速手段
• 批量推理：模型批量处理文本时能充分利用硬件并行能力，避免逐条推理的冗余开销
• 列表操作优化：用extend替代原代码的insert，insert在大列表中每次插入都要移动元素，时间复杂度为O(n)，而extend仅为O(k)（k为批量大小），效率大幅提升
• 冗余打印移除：删除循环内的逐行打印，改为批量进度提示，避免IO操作拖慢任务
• 断点续传机制：通过保存进度文件，中断后可从上次位置继续处理，无需重新运行全部60万+样本

额外建议 #

• 若内存不足，可通过pandas的chunksize参数分块读取数据集，避免一次性加载全部样本
• 提前过滤clean_text列的空值、异常文本，避免预测时因输入错误导致任务中断
• 若使用CPU环境，建议切换到云GPU服务（如Colab、国内云厂商GPU实例），否则即使优化后仍会耗时较长

内容的提问来源于stack exchange，提问作者Hiss