百万级嵌入矩阵两两余弦距离计算内存溢出问题求助

针对10万×100的嵌入矩阵计算两两余弦距离的内存溢出问题，给你几个可行的实操方案：

1. 计算后直接写入磁盘，不累积内存 #

你之前的分批计算崩溃，核心是把所有batch的结果都存在distances_matrix列表里，累计到11个batch时，已经占用了4.4GB内存，继续累加必然溢出。改成每算完一个batch就直接写入磁盘，用HDF5格式存储可以实现分块写入，内存里始终只保留当前batch的结果：

import h5py
import numpy as np
from sklearn.metrics.pairwise import cosine_distances
from tqdm import tqdm

embeddings = embeddings.astype(np.float32)
batch_size = 1000
df_size = len(embeddings)

# 创建HDF5文件并预分配存储空间
with h5py.File('cosine_distances.h5', 'w') as f:
    distances_dataset = f.create_dataset('distances', shape=(df_size, df_size), dtype=np.float32)
    
    for i in tqdm(range(0, df_size, batch_size)):
        end = min(i + batch_size, df_size)
        batch = embeddings[i:end]
        batch_distances = cosine_distances(batch, embeddings)
        # 将当前batch结果写入对应位置
        distances_dataset[i:end, :] = batch_distances

2. 手动基于归一化优化计算，减少内存开销 #

余弦距离公式为1 - (a·b)/(||a||*||b||)，先对所有嵌入做L2归一化后，||a||=||b||=1，余弦距离可简化为1 - a·b，计算更快且能减少中间变量的内存占用：

import numpy as np
from tqdm import tqdm
import h5py

# 对嵌入矩阵做L2归一化
embeddings_norm = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
embeddings_norm = embeddings_norm.astype(np.float32)

batch_size = 1000
df_size = len(embeddings_norm)

with h5py.File('cosine_distances_norm.h5', 'w') as f:
    distances_dataset = f.create_dataset('distances', shape=(df_size, df_size), dtype=np.float32)
    
    for i in tqdm(range(0, df_size, batch_size)):
        end = min(i + batch_size, df_size)
        batch_norm = embeddings_norm[i:end]
        # 矩阵点积得到余弦相似度，1减后为余弦距离
        batch_distances = 1 - batch_norm @ embeddings_norm.T
        distances_dataset[i:end, :] = batch_distances

手动用矩阵乘法计算比调用sklearn函数更省内存，避免了库内部的额外内存开销。

3. 降低数据精度（场景允许时） #

如果业务对精度要求不高，可以将数据从float32降到float16，直接减少一半的内存和磁盘占用：

embeddings_norm = embeddings_norm.astype(np.float16)

with h5py.File('cosine_distances_fp16.h5', 'w') as f:
    distances_dataset = f.create_dataset('distances', shape=(df_size, df_size), dtype=np.float16)
    # 后续计算逻辑同上，最终结果转float16写入

注意：float16精度有限，会存在微小误差，需先验证是否符合需求。

4. 用GPU加速计算（有GPU资源时） #

如果有GPU，用PyTorch或TensorFlow计算，GPU显存能更高效处理大矩阵，且计算速度更快：

import torch
from tqdm import tqdm

# 将嵌入矩阵转成GPU张量
embeddings_tensor = torch.tensor(embeddings, dtype=torch.float32).cuda()
# L2归一化
embeddings_tensor = torch.nn.functional.normalize(embeddings_tensor, p=2, dim=1)

batch_size = 2000  # 可根据GPU显存大小调整
df_size = len(embeddings_tensor)

with h5py.File('cosine_distances_gpu.h5', 'w') as f:
    distances_dataset = f.create_dataset('distances', shape=(df_size, df_size), dtype=np.float32)
    
    for i in tqdm(range(0, df_size, batch_size)):
        end = min(i + batch_size, df_size)
        batch_tensor = embeddings_tensor[i:end]
        # 计算余弦距离
        batch_distances = 1 - torch.matmul(batch_tensor, embeddings_tensor.T)
        # 转成numpy写入磁盘，同时释放GPU显存
        distances_dataset[i:end, :] = batch_distances.cpu().numpy()
        del batch_distances
        torch.cuda.empty_cache()

5. 只计算必要部分（无需完整矩阵时） #

如果最终目的不是获取完整距离矩阵，而是找每个样本的topK最近邻，直接用近似最近邻库（如faiss）检索，内存和计算量会大幅降低：

import faiss

# 归一化嵌入矩阵
embeddings_norm = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
embeddings_norm = embeddings_norm.astype(np.float32)

# 创建IVF_HNSW索引加速检索
index = faiss.IndexIVFHSW(faiss.IndexFlatIP(100), 100, 32)
index.train(embeddings_norm)
index.add(embeddings_norm)

# 查找每个样本的top10最近邻
k = 10
similarities, indices = index.search(embeddings_norm, k)
# 如需余弦距离，转换为：distances = 1 - similarities

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