文本分类中KNN元特征生成代码的优化需求求助

问题背景 #

我正在执行文本分类任务，需要从TF-IDF特征表示中提取KNN元特征，核心逻辑为：针对单个文档，从每个类别中选取k个最近邻文档（k=2、三类文档的生成示例如图）。当前代码在小数据集上运行正常，但在大数据集下计算复杂度显著升高，已多次优化到当前最优版本，恳请进一步优化指导。

原始代码 #

from scipy.spatial.distance import cdist
def scipy_cdist(point1, point2):
    distance = cdist(np.array([point1]), np.array([point2]), metric='euclidean')[0][0]
    return distance
# calculating the sorted distances for one document with others
def get_dist(train, test_row,y_train):
    distances = list()
    dist_ord=list()   
    for i in range(len(train)):
        dist = euclidean_distance(test_row, train[i])
        distances.append((train[i],y_train[i], dist))
    #sorted by classes then by distance
    distances.sort(key=itemgetter(1,2)) 
    return distances#,dist_ord

#calculating distances for all documents
all_distances_train=list()
for i in range(len(x_train)):
    d= get_dist(x_train, x_train[i],y_train)
    all_distances_train.append(d)      

#calculating a class neighbors(for one document)
def class_neighbors(dist,c,k):
    n=list()
    for i in range(len(dist)):
        if dist[i][1]==c and dist[i][2]!=0:
           n.append(dist[i][2])
    c_neighbors=n[:k] 
    return c_neighbors
#calculating all the classes  neighbors(for one document)
def get_all_neighbors(distances,y_data,k):
    classes=np.unique(y_data)
    all_neighbors=list()
    for c in range (len(classes)):
       n=class_neighbors(distances,c,k)
       all_neighbors.extend(n)
    return all_neighbors

#generating knnL2 metafeatures

def get_knnL2(all_distances,y_data,k):
    knnL2=list()
    for i in range(len(all_distances)):
        neighbors=get_all_neighbors(all_distances[i],y_data,k)
        knnL2.append(neighbors)
    return knnL2  

#generating knnL2_train metaFeatures
knnL2_train=get_knnL2(all_distances_train,y_train,2)
knnL2_train=np.array(knnL2_train)
#display  
print(knnL2_train)

#.......................knnl2_test
#calculating distances for all documents
all_distances_test=list()
for i in range(len(x_test)):
    d= get_dist(x_train, x_test[i],y_train)
    all_distances_test.append(d)      

#generating knnL2-test metafeatures
knnL2_test=get_knnL2(all_distances_test,y_train,2)
knnL2_test=np.array(knnL2_test)
#display
print(knnL2_test)

核心瓶颈分析 #

1. 逐样本距离计算效率极低：原始代码用Python循环逐个计算样本间欧氏距离，完全浪费了cdist的批量计算能力，Python循环在大数据集下开销极大。
2. 内存冗余严重：all_distances_train存储了所有文档对的完整元组（样本数据、类别、距离），大数据集下会直接导致内存溢出。
3. 排序与筛选步骤冗余：先全局按类别+距离排序，再逐个类别筛选k个非零距离样本，全局排序的时间复杂度极高，且不需要存储所有距离信息。

优化方案与代码实现 #

优化思路

• 用批量距离计算替代逐样本循环，利用numpy/scipy的向量化操作加速。
• 按类别分组处理样本，直接在每个类别组内筛选k个最近邻，避免全局排序。
• 仅存储必要的距离值，大幅降低内存占用。
• 利用TF-IDF的稀疏特性（可选），用稀疏矩阵专用的距离计算函数进一步提速。

优化后代码

import numpy as np
from scipy.spatial.distance import cdist
from sklearn.preprocessing import LabelEncoder

def get_knn_meta_features(X_train, y_train, X_test=None, k=2):
    # 类别编码，方便分组处理
    le = LabelEncoder()
    y_train_encoded = le.fit_transform(y_train)
    class_indices = [y_train_encoded == c for c in le.classes_]
    # 按类别分组训练样本（保留原始矩阵类型，支持稀疏矩阵）
    class_groups = [X_train[idx] for idx in class_indices]

    def process_target_samples(X_target):
        knn_features = []
        for sample in X_target:
            sample_2d = sample.reshape(1, -1)
            current_features = []
            for group in class_groups:
                # 批量计算当前样本与该类别所有样本的欧氏距离
                distances = cdist(sample_2d, group, metric='euclidean')[0]
                # 训练集处理时，排除样本自身（距离为0的情况）
                if np.array_equal(X_target, X_train):
                    self_idx = np.where((group == sample).all(axis=1))[0]
                    if self_idx.size > 0:
                        distances = np.delete(distances, self_idx)
                # 取k个最小距离，不足则用NaN填充（可根据需求调整填充策略）
                top_k_dist = np.sort(distances)[:k]
                if len(top_k_dist) < k:
                    top_k_dist = np.pad(top_k_dist, (0, k - len(top_k_dist)), 
                                       mode='constant', constant_values=np.nan)
                current_features.extend(top_k_dist.tolist())
            knn_features.append(current_features)
        return np.array(knn_features)

    # 生成训练集和测试集的KNN元特征
    knnL2_train = process_target_samples(X_train)
    knnL2_test = process_target_samples(X_test) if X_test is not None else None
    return knnL2_train, knnL2_test

# 使用示例（假设X_train/y_train/X_test为TF-IDF矩阵）
# knnL2_train, knnL2_test = get_knn_meta_features(X_train, y_train, X_test, k=2)
# print(knnL2_train)
# print(knnL2_test)

额外优化建议 #

• 近似最近邻算法：如果样本量超百万级，可使用FAISS、Annoy等库实现近似最近邻搜索，大幅降低计算时间（牺牲少量精度）。
• 并行计算：用joblib.Parallel并行处理每个样本的特征提取，利用多核CPU资源：

  from joblib import Parallel, delayed
  # 修改process_target_samples中的循环为并行
  def process_single_sample(sample, class_groups, is_train):
      sample_2d = sample.reshape(1, -1)
      current_features = []
      for group in class_groups:
          distances = cdist(sample_2d, group, metric='euclidean')[0]
          if is_train:
              self_idx = np.where((group == sample).all(axis=1))[0]
              if self_idx.size > 0:
                  distances = np.delete(distances, self_idx)
          top_k_dist = np.sort(distances)[:k]
          if len(top_k_dist) < k:
              top_k_dist = np.pad(top_k_dist, (0, k - len(top_k_dist)), 
                                 mode='constant', constant_values=np.nan)
          current_features.extend(top_k_dist.tolist())
      return current_features
  
  # 并行处理
  knn_features = Parallel(n_jobs=-1)(delayed(process_single_sample)(
      sample, class_groups, np.array_equal(X_target, X_train)) 
      for sample in X_target)
  

• 降维预处理：对高维TF-IDF矩阵先做TruncatedSVD降维，减少距离计算的维度，进一步提升速度。

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