为什么这种欧式距离计算方法会爆炸 RAM 使用量？答案

【问题标题】：Why does this euclidean distance calculation method explodes RAM usage?为什么这种欧式距离计算方法会爆炸 RAM 使用量？
【发布时间】：2021-10-22 07:46:36
【问题描述】：

我正在研究使用 2017 年斯坦福课程中的一些材料对图像进行分类的 KNN 算法。我们得到一个由许多图像组成的数据集，后来这些集合表示为 2D numpy 数组，我们应该编写计算这些图像之间距离的函数。更具体地说，给定测试图像的二维数组和训练图像的二维数组，我被要求编写一个 L_2 距离函数，它将这两个集合作为输入并返回一个距离矩阵，其中每一行 i 代表一个测试图像，每列 j 代表一个训练图像。

该练习还要求我不使用任何循环且不使用np.abs 函数。于是我试了试：

def compute_distances_no_loops(self, X):
    """
    Compute the distance between each test point in X and each training point
    in self.X_train using no explicit loops.

    Input / Output: Same as compute_distances_two_loops
    """
    num_test = X.shape[0]
    num_train = self.X_train.shape[0]
    dists = np.zeros((num_test, num_train))
    all_test_subs_sq = (X[:, np.newaxis] - self.X_train)**2
    dists = np.sqrt(np.sum(all_test_subs_sq), axis = 2)
    return dists

显然，由于分配了大约 60 GB 的 RAM，这会使 Google 的 Colab 环境在 6 秒内崩溃。我想我应该澄清一下训练集 X_train 的形状为 (5000, 3072)，而测试集 X 的形状为 (500, 3072)。我不确定这里会发生什么如此占用大量 RAM，但我又不是最聪明的人来计算空间复杂度。

我在 Google 上搜索了一下，发现了一个无需 NASA 计算机即可工作的解决方案，它使用平方和公式：

    dists = np.reshape(np.sum(X**2, axis=1), [num_test,1]) + np.sum(self.X_train**2, axis=1)\
- 2 * np.matmul(X, self.X_train.T)
    dists = np.sqrt(dists)

我也不确定为什么这个解决方案不像我的那样爆炸。非常感谢您在这里提供任何见解，非常感谢您的阅读。

【问题讨论】：

结果的大小是两个输入数组大小的乘积。

标签： python arrays numpy distance space-complexity

【解决方案1】：

在 compute_distances_no_loops() 函数中，中间数组 all_test_subs_sq 的形状为 (500, 3072, 5000)，因此它由 500 * 3072 * 5000 = 7,680,000,000 个元素组成。假设X和X_train的dtype为float64，每个元素的权重为8字节，那么数组的总大小为61,440,000,000字节，即约60GB。

您包含的另一个解决方案避免了这个问题，因为它不会创建如此大的中间数组。 np.reshape(np.sum(X**2, axis=1), [num_test,1]) 的形状是 (500, 1)，np.sum(self.X_train**2, axis=1) 的形状是 (5000,)。当您添加它们时，您将获得一个形状数组 (500, 5000)。 np.matmul(X, self.X_train.T) 也会产生一个相同形状的数组。

【讨论】：

【解决方案2】：

问题出在

all_test_subs_sq = (X[:, np.newaxis] - self.X_train)**2

X[:, np.newaxis] 等价于X[:, np.newaxis, :] 形状为(50, 1, 3072)。广播后，X[:, np.newaxis] - self.X_train 产生一个密集 (500, 5000, 3072) 数组，这是一个巨大的500 x 5000 x 3072 x 8 bytes ≈ 61.44 GB，因为你有np.float64。

【讨论】：