为什么使用 pytorch 广播的 KNN 这么慢？

问：

我正在尝试为网格点找到 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

我制作了一个自定义的 KNN 函数，以便它可以在 pytorch 中的 gpu 上运行。 pytorch 代码如下所示，假设 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 代码如下所示

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

我使用了广播，所以我希望它在 GPU 上运行得很快。但是，当我使用 .速度这么慢的原因是什么？timeit