如何计算超大数组下的scikit-learn rbf_kernel()并解决内存不足问题

sklearn自带的rbf_kernel默认会一次性计算全量样本的两两相似度，当样本量过大会生成超大规模的N×N核矩阵，直接触发内存溢出。我们可以通过分块计算的方式，无需一次性加载全量数据即可得到完整结果，方案如下：

适用场景1：全量特征可加载进内存，仅核矩阵计算时溢出 #

如果你的样本特征矩阵X可以完整存入内存，只是生成核矩阵时内存不足，可按行拆分计算：

from sklearn.metrics.pairwise import rbf_kernel
import numpy as np

# 你的全量特征矩阵，形状为 (n_samples, n_features)
X = np.random.randn(10000, 128)  # 示例数据，可替换为你的实际数据
gamma = 60
batch_size = 1000  # 可根据可用内存调整，越大计算效率越高
n_samples = X.shape[0]

# 若需要完整存储核矩阵，先初始化空矩阵（不需要存储可跳过，边算边处理结果）
kernel_matrix = np.zeros((n_samples, n_samples), dtype=np.float32)  # 用float32比默认float64省一半内存

for i in range(0, n_samples, batch_size):
    # 取当前批次的特征
    X_batch = X[i:i+batch_size]
    # 计算当前批次和全量样本的RBF核结果，形状为 (batch_size, n_samples)
    batch_res = rbf_kernel(X_batch, X, gamma=gamma)
    # 写入对应位置，不需要存全量的话此处直接处理batch_res即可
    kernel_matrix[i:i+batch_size] = batch_res

适用场景2：连特征矩阵都无法全量加载进内存 #

如果样本量极大，特征矩阵本身就无法存入内存，可采用双重分块+磁盘存储的方案，每次只加载两个小批次的特征计算子核矩阵，最终拼接为完整结果：

from sklearn.metrics.pairwise import rbf_kernel
import numpy as np
import h5py  # 用HDF5格式存储大文件，无需全量加载进内存

# 假设你的全量特征存在HDF5文件中，key为X
with h5py.File("your_large_features.h5", "r") as f:
    n_samples = f["X"].shape[0]
    gamma = 60
    batch_size = 500  # 批次调得越小，内存占用越低
    
    # 核矩阵也直接写入HDF5文件，不占用内存
    with h5py.File("rbf_kernel_result.h5", "w") as kf:
        kernel_dset = kf.create_dataset("kernel", (n_samples, n_samples), dtype=np.float32)
        
        # 遍历左侧批次
        for i in range(0, n_samples, batch_size):
            X_left = f["X"][i:i+batch_size]
            # 遍历右侧批次
            for j in range(0, n_samples, batch_size):
                X_right = f["X"][j:j+batch_size]
                # 计算两个批次间的RBF核
                sub_kernel = rbf_kernel(X_left, X_right, gamma=gamma)
                # 写入对应位置
                kernel_dset[i:i+batch_size, j:j+batch_size] = sub_kernel

优化建议 #

• 精度兼容优化：大部分场景下用float32存储核结果即可满足需求，相比默认float64直接减少一半内存占用
• 批次大小调整：根据可用内存灵活调整batch_size，平衡计算效率和内存占用
• 按需处理结果：如果不需要完整的核矩阵（比如仅需要取每行TopK相似样本），可以每计算完一个批次就直接处理结果，无需存储整个核矩阵，进一步降低内存消耗

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