如何优化NLP项目中文本特征提取的运行速度与内存占用

现有代码核心瓶颈 #

• Pandas逐行apply本质是Python层循环，10万级样本下执行开销极高
• 求共同词时使用列表的in查询，时间复杂度为O(n)，重复计算开销大
• 额外存储两份分词后的列表列，Python字符串对象内存占用远高于向量化存储结构

可落地的优化方案 #

1. 低改造成本优化（收益提升3-10倍） #

• 成员查询替换为集合操作：将分词后的列表转为集合，in查询复杂度降到O(1)，共同词计算直接取集合交集长度
• 用Pandas内置向量化函数替代逐行计算：比如长度类特征直接用.str.len()向量化计算，无需走逐行逻辑
• 特征一次性生成：让特征函数直接返回Series，避免后续多次map拆分特征

2. 中等改造成本优化（收益提升10-100倍） #

• 用Numba JIT编译复杂特征逻辑：给纯Python实现的特征函数加@numba.jit(nopython=True)装饰器，直接编译为机器码执行，避开Python解释器开销
• 批量处理NLP前置步骤：用spaCy的nlp.pipe批量处理分词、NER等操作，比逐行调用nltk效率高5-20倍

3. 内存优化方案 #

• 不需要留存的中间结果（比如分词列表）计算完特征后直接删除，避免占用内存
• 特征计算完成后统一转为numpy数组存储，比Pandas的object类型列内存占用低60%以上
• 超大数据量下可分块处理，每处理1-2万条样本就释放中间变量内存，再处理下一批

优化后代码示例 #

import nltk
import numpy as np
import pandas as pd
from functools import partial

# 示例数据
df = pd.DataFrame([["The quick brown fox jumps over the lazy dog.",
                    "Energy is sustainable if it meets the needs of the present without compromising the ability of future generations to meet their needs."],
                   ["The scientific literature on limiting global warming describes pathways in which the world rapidly phases out coal-fired power plants, produces more electricity from clean sources such as wind and solar, shifts towards using electricity instead of fuels in sectors such as transport and heating buildings, and takes measures to conserve energy.",
                    "Lorem Ipsum is simply dummy text of the printing and typesetting industry. Lorem Ipsum has been the industry's standard dummy text ever since the 1500s"]], columns=['text1', 'text2'])

def process(text):
    tokens = nltk.word_tokenize(text)
    # 其他预处理如 stemming、lemmatization
    # 直接返回集合，省得后续重复转
    return set(tokens), len(tokens)

# 批量处理分词，同时拿到长度和token集合，避免存两份大列
process_partial = partial(process)
df[['text1_set', 'text1_len']] = df['text1'].apply(process_partial).apply(pd.Series)
df[['text2_set', 'text2_len']] = df['text2'].apply(process_partial).apply(pd.Series)

# 向量化计算特征，完全避免逐行apply
df['feature1'] = df['text1_len'] + df['text2_len']
# 共同词长度直接用集合交集
df['feature2'] = df.apply(lambda x: len(x['text1_set'] & x['text2_set']), axis=1)

# 用完删除中间结果释放内存
df.drop(columns=['text1_set', 'text2_set', 'text1_len', 'text2_len'], inplace=True)

numpy向量化落地说明 #

对于数值类特征，可以将所有特征列转为numpy矩阵统一运算，比如：

# 将所有特征转为numpy数组，后续运算完全在numpy层执行
feature_matrix = df[['feature1', 'feature2']].values
# 后续特征变换、归一化等操作都可以用numpy向量化函数完成，速度远快于Python循环

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