你不太可能在不仔细考虑的情况下完整地计算出一个1M*1M的矩阵。您可能需要沿着scipy.spatial.KDTree行的内容。一旦构建了一棵树，就可以将代理的坐标传递给query方法，以获取特定半径内的邻居。为了一次得到所有的邻居，你可以计算树的sparse_distance_matrix，它本身在一个适当的阈值。
或者，您可以研究任意数量的高效聚类算法。

我找到了三个解决方案，解决方案1

import torch
x = torch.randn(3000000, 2).cuda()
y = x

# Turn our Tensors into KeOps symbolic variables:
from pykeops.torch import LazyTensor
x_i = LazyTensor( x[:,None,:] ) 
y_j = LazyTensor( y[None,:,:] )  

# We can now perform large-scale computations, without memory overflows:
D_ij = ((x_i - y_j)**2).sum(dim=2)  
D_ij.argKmin(20,dim=1)

溶液2

M = 3000000
import numpy as np
from pykeops.numpy import LazyTensor as LazyTensor_np
x = np.random.rand(M, 2)
y = x
x_i = LazyTensor_np(
    x[:, None, :]
)  # (M, 1, 2) KeOps LazyTensor, wrapped around the numpy array x
y_j = LazyTensor_np(
    y[None, :, :]
)  # (1, N, 2) KeOps LazyTensor, wrapped around the numpy array y

D_ij = ((x_i - y_j) ** 2).sum(-1)  # **Symbolic** (M, N) matrix of squared distances
s_i = D_ij.argKmin(20,dim=1).ravel()  # genuine (M,) array of integer indices

溶液3

from sklearn.neighbors import NearestNeighbors
import numpy as np
M = 3000000
x = np.random.rand(M, 2)
nbrs = NearestNeighbors(n_neighbors=20, algorithm='ball_tree').fit(x)
distances, indices = nbrs.kneighbors(x)

虽然三种解决方案的执行时间都是一样的，一分钟，但对内存的需求分别约为2GB、1GB和1.3GB。这将是伟大的听到的想法，以降低执行时间。

一个简单的解决方案，没有任何花哨的库：

def batch_nn(X, Y, bs):
    nns = torch.zeros(len(X))
    for i in range(0, len(X), bs):
        batch_x = X[i:i+bs]
        dists = torch.cdist(batch_x, Y)
        nns[i:i+bs] = torch.argmin(dists, axis=1)
    return nns

然后运行batch_nn(X, X, 10)。
这在我的笔记本电脑上大约需要一分钟，你提到的X。但它不会坠毁。这是很容易扩展到寻找更多的比最近的邻居以及。使用torch.argpartition或torch.topk代替torch.argmin。