训练Multinomial Naive Bayes模型时遇内存错误，求优化方案

这问题我之前做大规模文本分类时也碰到过，核心原因就是你的特征提取逻辑太消耗内存了——每个评论都生成一个包含所有候选词的字典，10万条数据堆下来，32GB内存根本扛不住。给你几个亲测有效的解决思路：

1. 用稀疏矩阵替代稠密字典 #

你的features[word] = (word in set(comment))本质是在生成稠密布尔特征，但绝大多数词在单个评论里都是不存在的，这种存储方式完全是浪费内存。改用稀疏矩阵只存储存在的词的信息，内存占用能降到原来的几十分之一甚至几百分之一。

举个具体的实现例子：

from collections import Counter
import numpy as np
from scipy.sparse import csr_matrix

# 第一步：先统计全量数据的词频，过滤低频词（减少特征维度）
all_words = []
for comment in full_dataset_comments:
    all_words.extend(comment.split())  # 假设你的评论已经分词，若未分词先做分词处理
word_counter = Counter(all_words)
# 只保留出现次数≥5的词（阈值可根据数据调整）
valid_words = [word for word, cnt in word_counter.items() if cnt >= 5]
word_to_idx = {word: idx for idx, word in enumerate(valid_words)}
num_features = len(valid_words)

# 第二步：构建稀疏特征矩阵
def build_sparse_features(comments):
    rows, cols, data = [], [], []
    for sample_idx, comment in enumerate(comments):
        comment_words = set(comment.split())
        for word in comment_words:
            if word in word_to_idx:
                rows.append(sample_idx)
                cols.append(word_to_idx[word])
                data.append(1)  # 表示该词存在
    # 用uint8存储布尔值，进一步节省内存
    return csr_matrix((data, (rows, cols)), shape=(len(comments), num_features), dtype=np.uint8)

这样生成的csr_matrix只会存储每个评论中实际出现的词，内存占用会大幅降低。

2. 批量处理+增量训练 #

不要一次性加载并处理10万条数据，拆成小批量（比如每1000条一批），处理一批就训练一批，训练完就释放该批次的内存，再处理下一批。

示例代码：

batch_size = 1000
total_samples = len(full_dataset_comments)
num_batches = (total_samples + batch_size - 1) // batch_size  # 向上取整

# 假设你的模型支持增量训练（比如sklearn的SGDClassifier）
from sklearn.linear_model import SGDClassifier
model = SGDClassifier(loss='log_loss')

for batch_idx in range(num_batches):
    start = batch_idx * batch_size
    end = min(start + batch_size, total_samples)
    batch_comments = full_dataset_comments[start:end]
    batch_labels = full_dataset_labels[start:end]
    
    # 构建当前批次的稀疏特征
    batch_features = build_sparse_features(batch_comments)
    
    # 增量训练模型
    model.partial_fit(batch_features, batch_labels, classes=np.unique(full_dataset_labels))
    
    # 手动释放内存
    del batch_features
    import gc
    gc.collect()

这种方式能把内存占用控制在单个批次的量级，32GB内存完全够用。

3. 优化特征提取的细节 #

• 提前过滤低频词：上面的例子里已经用到了，低频词不仅对模型贡献小，还会大幅增加特征维度，过滤掉能直接减少内存压力。
• 避免重复创建Set：如果你的评论是长文本，每次都set(comment)会重复计算，可以考虑先对所有评论做一次分词并转成集合存储，避免重复操作。
• 用生成器代替列表：如果你的get_features是一次性生成所有评论的特征列表，改成生成器（用yield代替return），每次只生成一个评论的特征，不会一次性把所有特征加载到内存。

4. 利用外部存储（极端情况） #

如果处理完上面的步骤后内存还是紧张，可以把特征矩阵存储到磁盘上（比如用HDF5格式），训练时再分批读取：

import h5py

# 分批写入特征到HDF5文件
with h5py.File('features.h5', 'w') as f:
    dset = f.create_dataset('features', shape=(total_samples, num_features), dtype='uint8', compression='gzip')
    for batch_idx in range(num_batches):
        start = batch_idx * batch_size
        end = min(start + batch_size, total_samples)
        batch_features = build_sparse_features(full_dataset_comments[start:end])
        dset[start:end] = batch_features.toarray()  # 若稀疏矩阵太大，也可以分批存储稀疏格式

# 训练时分批读取
with h5py.File('features.h5', 'r') as f:
    features = f['features']
    for batch_idx in range(num_batches):
        start = batch_idx * batch_size
        end = min(start + batch_size, total_samples)
        batch_features = features[start:end]
        # 训练模型...

按照这些方法调整后，别说10万条，50万条数据也能轻松处理。

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