使用Universal Sentence Encoder处理60M行数据内存不足的解决方案咨询

针对6000万行文本生成的512维嵌入向量（约120GB浮点数据），直接存储为numpy数组内存不足的问题，推荐以下几种高效方案：

1. 分块生成+增量存储 #

不要一次性生成所有嵌入，而是分批处理文本数据，每生成一批就写入磁盘，避免全量数据占满内存。

示例代码（TensorFlow + h5py） #

import tensorflow_hub as hub
import h5py
import pandas as pd

# 加载Universal Sentence Encoder
encoder = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")

# 分块读取文本数据（假设存储在csv文件中）
chunk_size = 100000
total_rows = 60000000

# 创建HDF5文件存储嵌入
with h5py.File('embeddings.h5', 'w') as f:
    # 创建可扩展的数据集，指定分块大小
    emb_dataset = f.create_dataset(
        'embeddings',
        shape=(0, 512),
        maxshape=(total_rows, 512),
        dtype='float32',
        chunks=(chunk_size, 512)
    )
    
    # 分块处理文本
    for i, chunk in enumerate(pd.read_csv('text_data.csv', chunksize=chunk_size)):
        texts = chunk['text_column'].tolist()
        # 生成嵌入
        embeddings = encoder(texts).numpy()
        # 写入HDF5
        start_idx = i * chunk_size
        end_idx = start_idx + len(embeddings)
        emb_dataset.resize(end_idx, axis=0)
        emb_dataset[start_idx:end_idx] = embeddings
        print(f"完成第{i+1}批嵌入存储，已处理{end_idx}条数据")

2. 用高效存储格式替代numpy数组 #

.npy格式需要一次性加载全量数据，换成HDF5或Parquet这类支持分块读写的格式，后续使用时可以按需加载部分数据，不用占满内存。

Parquet格式示例（pandas） #

# 分块生成嵌入并写入Parquet
chunk_size = 100000
for i, chunk in enumerate(pd.read_csv('text_data.csv', chunksize=chunk_size)):
    texts = chunk['text_column'].tolist()
    embeddings = encoder(texts).numpy()
    # 转为DataFrame写入Parquet
    emb_df = pd.DataFrame(embeddings, columns=[f'emb_{j}' for j in range(512)])
    # 追加模式写入
    emb_df.to_parquet(
        'embeddings.parquet',
        engine='pyarrow',
        append=True,
        compression='snappy'
    )

# 后续读取时，同样可以分块加载
for emb_chunk in pd.read_parquet('embeddings.parquet', chunksize=chunk_size):
    # 处理当前块的嵌入
    pass

3. 降低嵌入向量维度 #

通过降维算法压缩嵌入维度，减少内存占用。比如用PCA把512维降到128维，内存需求直接降到原来的1/4（约30GB），如果继续降到64维则仅15GB，大部分服务器内存都能容纳。

PCA降维示例 #

from sklearn.decomposition import PCA

# 先生成一小批嵌入训练PCA模型
sample_texts = pd.read_csv('text_data.csv', nrows=10000)['text_column'].tolist()
sample_emb = encoder(sample_texts).numpy()
pca = PCA(n_components=128)
pca.fit(sample_emb)

# 分块生成嵌入并降维后存储
with h5py.File('embeddings_pca.h5', 'w') as f:
    emb_dataset = f.create_dataset(
        'embeddings',
        shape=(0, 128),
        maxshape=(total_rows, 128),
        dtype='float32',
        chunks=(chunk_size, 128)
    )
    
    for i, chunk in enumerate(pd.read_csv('text_data.csv', chunksize=chunk_size)):
        texts = chunk['text_column'].tolist()
        embeddings = encoder(texts).numpy()
        # 降维
        embeddings_pca = pca.transform(embeddings)
        # 写入
        start_idx = i * chunk_size
        end_idx = start_idx + len(embeddings_pca)
        emb_dataset.resize(end_idx, axis=0)
        emb_dataset[start_idx:end_idx] = embeddings_pca

4. 使用numpy内存映射数组 #

numpy的memmap可以把磁盘文件当作内存数组使用，数据不会一次性加载到内存，而是按需读取/写入，适合超大数组的存储。

内存映射示例 #

import numpy as np

# 创建内存映射文件，初始大小设为总数据量
emb_shape = (total_rows, 512)
memmap_arr = np.memmap(
    'embeddings_memmap.npy',
    dtype='float32',
    mode='w+',
    shape=emb_shape
)

# 分块写入
for i, chunk in enumerate(pd.read_csv('text_data.csv', chunksize=chunk_size)):
    texts = chunk['text_column'].tolist()
    embeddings = encoder(texts).numpy()
    start_idx = i * chunk_size
    end_idx = start_idx + len(embeddings)
    memmap_arr[start_idx:end_idx] = embeddings
    # 强制写入磁盘
    memmap_arr.flush()

# 后续读取时，同样用memmap加载，不会占满内存
loaded_memmap = np.memmap('embeddings_memmap.npy', dtype='float32', mode='r', shape=emb_shape)
# 按需读取部分数据
batch = loaded_memmap[0:10000]

5. 直接增量训练分类模型 #

如果最终目标是分类任务，完全不用先存储所有嵌入。可以生成一批嵌入就训练一批模型，用增量学习的方式完成训练，全程不用保存全量嵌入。

增量训练示例（scikit-learn） #

from sklearn.linear_model import SGDClassifier
from sklearn.metrics import accuracy_score

# 初始化增量分类器
clf = SGDClassifier(loss='log_loss')

# 分块处理数据，假设标签列是'label'
for chunk in pd.read_csv('text_data.csv', chunksize=chunk_size):
    texts = chunk['text_column'].tolist()
    labels = chunk['label'].values
    embeddings = encoder(texts).numpy()
    # 增量训练
    clf.partial_fit(embeddings, labels, classes=np.unique(labels))
    # 可选：验证当前批次
    preds = clf.predict(embeddings)
    print(f"当前批次准确率: {accuracy_score(labels, preds):.4f}")

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