基于PyTorch广播实现的GPU版KNN为何远慢于Sklearn CPU版？

我为网格点实现KNN算法，生成网格的代码如下：

def grid_by(lims=[[0, 1], [0, 1]], size=[28, 28]):
    """
    Creates a tensor of 2D grid points.
    Grid points have one-to-one correspondence with input pixel values that are flattened in row-major order.

    Args:
        lims: [[domain of H], [domain of W]]
        size: [H, W]
    Returns:
        grid: Tensor of shape [(H*W), 2]
    """
    assert len(size) == 2 and len(lims) == len(size)
    expansions = [torch.linspace(start, end, steps) if i != 0 else torch.linspace(end, start, steps) for i, ((start, end), steps) in enumerate(zip(lims, size))]
    grid = torch.index_select(torch.cartesian_prod(*expansions),
                        dim=1,
                        index=torch.tensor([1,0]))
    return grid

基于PyTorch实现的GPU版KNN（暴力搜索+L2距离）代码：

def knn(grid, k):
    """
    Brute Force KNN.

    Args:
            grid: Tensor of shape [(H*W), D]
            k: Int representing number of neighbors 
    """
    d = grid.shape[-1]
    Xr = grid.unsqueeze(1)
    Yr = grid.view(1, -1, d)
    distances = torch.sqrt(torch.sum((Xr - Yr)**2, -1))
    dist, index = distances.topk(k, largest=False, dim=-1)
    return dist, index

grid = grid_by().to('cuda')
knn_dist, knn_index = knn(grid, k=10)

Sklearn的CPU版实现：

grid = grid_by()
nn = NearestNeighbors(n_jobs=-1)
nn.fit(grid)
knn_dist, knn_index = nn.kneighbors(grid, n_neighbors=10)

原以为PyTorch GPU版会更快，但测试后发现速度远慢于Sklearn CPU版，原因如下：

• 算法类型差异：Sklearn的NearestNeighbors默认不会用暴力搜索，它会根据数据规模自动选择KD-Tree或Ball-Tree这类近似近邻算法，时间复杂度是O(N logN)；而你的PyTorch实现是纯暴力的O(N²)计算，哪怕有GPU并行，当N（28×28=784）较小时，暴力搜索的开销反而比树结构更高。

• GPU的额外开销抵消优势：你的网格数据量极小，GPU处理这种小规模数据时，内核启动延迟、CPU-GPU数据传输的开销会远大于并行计算带来的收益。而Sklearn在CPU上直接处理数据，没有跨设备传输的额外成本，还能充分利用CPU缓存。

• 底层优化程度差距：Sklearn的近邻算法底层是高度优化的C/C++代码，用到了MKL、OpenBLAS等数学库的向量指令（SSE、AVX）优化，再加上n_jobs=-1启用多线程，效率极高；而你的PyTorch代码虽然用了broadcasting，但暴力计算产生的中间张量（比如784×784×2的差值张量）会占用GPU内存，且sqrt、sum等操作没有针对这种特定场景做定制优化，整体效率不如Sklearn的底层实现。

• GPU的并行特性不匹配小规模数据：GPU擅长处理大规模数据的并行计算，当数据量小时，GPU的大量核心无法被充分利用，线程调度、内存访问的额外开销会让它的表现不如CPU。而CPU在小规模数据上，凭借低延迟的缓存和高效的单/多线程执行，更有优势。

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