使用gensim训练千万级文档Word2Vec模型崩溃如何优化？

问题根因 #

你当前遇到的崩溃是因为df['text'].str.split(r'[\s]+')会一次性将1000万份文档的分词结果全部加载到内存，中间生成的巨型列表会直接耗尽内存资源，触发会话崩溃。

最优实现方案 #

Gensim的Word2Vec原生支持流式输入，只要传入可迭代对象即可，无需一次性把全量分词结果加载到内存，推荐使用可复用的语料迭代器实现：

1. 基于内存中DataFrame的流式迭代实现 #

定义可重复迭代的语料类，每次遍历都逐行处理文本，不会存储全量分词结果：

from gensim.models import Word2Vec
import re

class PandasCorpus:
    def __init__(self, df):
        self.df = df
        self.split_pattern = re.compile(r'[\s]+')
    def __iter__(self):
        # 逐行遍历df的text列，实时分词返回
        for text in self.df['text']:
            yield self.split_pattern.split(text.strip())

# 初始化语料迭代器
corpus = PandasCorpus(df)
# 直接传入迭代器训练即可，workers参数设为CPU核心数可大幅提升训练速度
model = Word2Vec(sentences=corpus, min_count=50, workers=8)

2. 超大数据量优化方案（无需全量数据加载到内存） #

如果后续数据规模进一步扩大无法全部放进内存，可以提前将分词后的文本逐行写入本地文本文件，每行对应一份文档的分词结果（用空格分隔），再通过文件迭代器训练，内存占用可以降到MB级：

class FileCorpus:
    def __init__(self, file_path):
        self.file_path = file_path
    def __iter__(self):
        with open(self.file_path, 'r', encoding='utf-8') as f:
            for line in f:
                yield line.strip().split()

corpus = FileCorpus("你的分词结果文件路径.txt")
model = Word2Vec(sentences=corpus, min_count=50, workers=8)

可选的分步训练写法 #

如果需要单独控制词汇构建和训练流程，也可以拆分为两步执行：

corpus = PandasCorpus(df)
model = Word2Vec(min_count=50, workers=8)
# 第一步：统计词汇表
model.build_vocab(corpus)
# 第二步：训练模型
model.train(corpus, total_examples=model.corpus_count, epochs=model.epochs)

内容的提问来源于stack exchange，提问作者ℕʘʘḆḽḘ