为何双循环实现KNN曼哈顿距离计算比向量化方法更快？

问题背景 #

我正在学习斯坦福大学UFLDL深度学习课程，作业要求实现KNN算法，需构建距离矩阵D，其中D[i][j]代表训练集第i个样本与测试集第j个样本的曼哈顿距离（L1范数）。

我实现了三种计算方法，并通过%timeit测试性能，结果发现通常不被推荐的Python双循环反而最快，想知道是实现错误还是该场景下循环更具优势。

三种实现方法 #

1. 双循环方法 #

def compute_distances_two_loops(train_X, test_X):
    '''
    Computes L1 distance from every sample of X to every training sample
    Uses simplest implementation with 2 Python loops

    Arguments:
    train_X, np array (num_train_samples, num_features)
    test_X,  np array (num_test_samples, num_features)


    Returns:
    dists, np array (num_test_samples, num_train_samples) - array
       with distances between each test and each train sample
    '''

    num_train = train_X.shape[0]
    num_test = test_X.shape[0]
    dists = np.zeros((num_test, num_train), np.float32)
    for i_test in range(num_test):
        for i_train in range(num_train):
            dists[i_test][i_train] = np.sum(np.abs(test_X[i_test] - train_X[i_train]))
 
    return dists

2. 单循环方法 #

def compute_distances_one_loop(train_X, test_X):
    '''
    Arguments:
    train_X, np array (num_train_samples, num_features)
    test_X,  np array (num_test_samples, num_features)
    
    Returns:
    dists, np array (num_test_samples, num_train_samples) - array
       with distances between each test and each train sample
    '''

    num_train = train_X.shape[0]
    num_test = test_X.shape[0]
    dists = np.zeros((num_test, num_train), np.float32)
    for i_test in range(num_test):
        diff_matrix = np.abs(train_X - np.tile(test_X[i_test], (num_train, 1)))
        to_sum_matrix = np.ones((train_X.shape[1], 1), np.float32)
        dists[i_test] = (diff_matrix @ to_sum_matrix).T

    return dists

3. 无循环全向量化方法 #

def compute_distances_no_loops(train_X, test_X):
    '''
    Computes L1 distance from every sample of X to every training sample
    Fully vectorizes the calculations using numpy

    Arguments:
    train_X, np array (num_train_samples, num_features)
    test_X,  np array (num_test_samples, num_features)
    
    Returns:
    dists, np array (num_test_samples, num_train_samples) - array
       with distances between each test and each train sample
    '''

    num_train = train_X.shape[0]
    num_test = test_X.shape[0]
   
    large_row_matrix_test = np.repeat(test_X, num_train, axis=0)
    large_row_matrix_train = np.tile(train_X, (num_test, 1))

    diff_matrix = np.abs(large_row_matrix_test - large_row_matrix_train)
    to_sum_matrix = np.ones((train_X.shape[1], 1), np.float32)

    dists = np.reshape(diff_matrix @ to_sum_matrix, (num_test, num_train))

    return dists

性能测试结果 #

• 双循环：614 ms ± 27.9 ms per loop（7次运行，每次1循环的均值±标准差）
• 单循环：836 ms ± 35.1 ms per loop
• 无循环：961 ms ± 30.5 ms per loop

解惑分析 #

你的双循环更快，核心原因是后两种向量化实现存在不必要的内存开销和低效操作，并非Python循环本身在这个场景更优：

1. 无循环方法的问题：

   • 通过np.repeat和np.tile生成的两个巨大中间矩阵，形状为(num_test*num_train, num_features)，会触发大量内存分配和数据复制，这部分开销远超过循环的开销。
   • 使用矩阵乘法@求和完全没必要，np.sum(diff_matrix, axis=1)直接按行求和效率高得多，矩阵乘法会引入额外计算步骤。

2. 单循环方法的问题：

   • 每次循环里的np.tile(test_X[i_test], (num_train, 1))会生成(num_train, num_features)的临时矩阵，同样存在内存复制开销。
   • 用矩阵乘法求和属于低效操作，替换成np.sum(diff_matrix, axis=1)能大幅提速。

3. 双循环的隐性优势：

   • 每次计算只处理两个向量的差与求和，内存占用极小，无多余中间矩阵。
   • 虽然是Python循环，但内部的np.sum和np.abs都是numpy的C实现，核心计算逻辑由高效底层代码完成，并非纯Python解释器执行。

优化后的向量化实现 #

正确的全向量化L1距离计算应利用numpy广播机制，避免生成超大中间矩阵：

def compute_distances_vectorized(train_X, test_X):
    # 利用广播，test_X扩展为(num_test, 1, num_features)，train_X为(1, num_train, num_features)
    diff = np.abs(test_X[:, np.newaxis, :] - train_X[np.newaxis, :, :])
    dists = np.sum(diff, axis=2)
    return dists

该实现无循环且无冗余内存开销，性能会远优于你之前的三种方法，可测试对比。

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