如何在 Python 中加速嵌套？

Question

我有一个包含描述符（数组）的列表，需要计算它们之间的余弦距离，这需要大量计算。 对于具有n向量的列表，它执行n * (n - 1) / 2操作。 我需要在一个大的n值中做到这一点，那么，如何使用多处理来加速这个过程？ 我找到了一些答案，但我不清楚在我的情况下该怎么做。

我想要加速的代码是下面的代码：

from scipy.spatial.distance import cosine as cosine_distance

n = len(all_vectors)
all_distances = []
for i in range(n):
    for j in range(i + 1, n):
        x1 = all_vectors[i]
        x2 = all_vectors[j]
        all_distances.append(cosine_distance(x1, x2))

[更新]

我还需要做一点标签检查，所以，原始代码如下：

from scipy.spatial.distance import cosine as cosine_distance

n = len(all_vectors)
all_distances = []
all_label_checks = []
for i in range(n):
    for j in range(i + 1, n):
        x1 = all_vectors[i]
        x2 = all_vectors[j]
        all_distances.append(cosine_distance(x1, x2))
        
        label_x1 = all_labels[i]  # all_labels and all_distances are tied
        label_x2 = all_labels[j]
        all_label_checks.append(int(label_x1 == label_x2))

我测试了一些建议的时间，到目前为止最好的答案来自@DaniMesejo。 这是我用来测试每​​个案例的代码：

import time
import numpy as np

from scipy.spatial.distance import pdist, cosine as cosine_distance, squareform
from itertools import combinations, starmap


EMBEDDINGS_SIZE = 128
N_EMBEDDINGS = 1024


def get_n_embeddings(n):
    return np.random.randn(n, EMBEDDINGS_SIZE)


embeddings = get_n_embeddings(N_EMBEDDINGS)

###################
# scipy pdist
###################
start = time.time()
pdist_distances = pdist(embeddings, 'cosine')
end = time.time()
print(f"Using scipy pdsit: {end - start}")

###################
# nested loop
###################
nested_distances = []
start = time.time()
for i in range(N_EMBEDDINGS):
    for j in range(i + 1, N_EMBEDDINGS):
        x1 = embeddings[i]
        x2 = embeddings[j]
        nested_distances.append(cosine_distance(x1, x2))
end = time.time()
print(f"Using nested loop: {end - start}")

###################
# combinations
###################
start = time.time()
combination_distances = []
for x1, x2 in combinations(embeddings, 2):
    combination_distances.append(cosine_distance(x1, x2))
end = time.time()
print(f"Using combination: {end - start}")

###################
# starmap
###################
start = time.time()
starmap_distances = list(starmap(cosine_distance, combinations(embeddings, 2)))
end = time.time()
print(f"Using starmap: {end - start}")

print(np.allclose(pdist_distances, nested_distances))
print(np.allclose(pdist_distances, pdist_class_distances))
print(np.allclose(pdist_distances, combination_distances))
print(np.allclose(pdist_distances, starmap_distances))

结果：

Using scipy pdsit: 0.0303647518157959
Using pdsit and class: 0.09841656684875488
Using nested loop: 13.415924549102783
Using combination: 13.093504428863525
Using starmap: 13.177483081817627
True
True
True
True

为了解决我的标签问题，我也可以使用pdist 。 我只需要在 2d 列表中转换我的标签列表（它告诉每个嵌入的标签是什么）并计算成对距离。 相同标签对的距离为 0：

###################
# pdist and class
###################
classes = []
count = 0
id_ = 0
for i in range(N_EMBEDDINGS):
    classes.append(id_)
    count += 1
    if count % 4 == 0:
        count = 0
        id_ += 1
labels_x = [(i, i) for i in classes]
start = time.time()
pdist_class_distances = pdist(embeddings, 'cosine')
pdist_labels = pdist(labels_x, 'euclidean')
pdist_labels = [1 if x == 0.0 else 0 for x in pdist_labels]
end = time.time()
print(f"Using pdsit and class: {end - start}")

Answer 1

一种可能的方法是使用 scipy 函数pdist ，例如：

from scipy.spatial.distance import pdist, cosine as cosine_distance, squareform
import numpy as np


def all_pairs_using_pdist(X):
    """Function using pdist"""
    distances = squareform(pdist(X, 'cosine'))
    rows, _ = distances.shape
    return list(distances[np.triu_indices(rows, 1)])


def all_pairs_for_loops(all_vectors):
    """Function using a nested for loop"""
    n = len(all_vectors)
    all_distances = []
    for i in range(n):
        for j in range(i + 1, n):
            x1 = all_vectors[i]
            x2 = all_vectors[j]
            all_distances.append(cosine_distance(x1, x2))
    return all_distances

对于如下矩阵：

Y = np.random.randn(100, 25)

它给出了时间：

%timeit all_pairs_using_pdist(Y)
630 µs ± 28.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit all_pairs_for_loops(Y)
216 ms ± 78 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

大约快300 倍。 当然，这两种方法都会产生相同的结果：

Y = np.random.randn(100, 25)
print(np.allclose(all_pairs_for_loops(Y), all_pairs_using_pdist(Y)))

输出

True

Answer 2

我不能说这有多大帮助，但是您的嵌套循环基本上是尝试使所有唯一的合法索引的 2 组合得到值的所有唯一 2 组合。 但是有内置函数可以做到这一点，所以使用from itertools import combinations ，这个：

for i in range(n):
    for j in range(i + 1, n):
        x1 = all_vectors[i]
        x2 = all_vectors[j]

可以简化为：

for x1, x2 in combinations(all_vectors, 2):

您甚至可以将所有这些放在itertools.starmap以将整个循环推送到 C 层（假设cosine_distance本身是在 C 中实现的；如果不是，它只会限制上行，但仍然可以正常工作），将您发布的所有代码简化为：

from scipy.spatial.distance import cosine as cosine_distance
from itertools import combinations, starmap

all_distances = list(starmap(cosine_distance, combinations(all_vectors, 2)))

并且运行速度更快（如果cosine_distance是在 C 中实现的成本相对较低的操作，则cosine_distance更快）。

如果这还不够快，很可能你将不得不使用multiprocessing / concurrent.futures.ProcessPoolExecutor来并行化（如果scipy函数是在 C 中实现的，并且在运行时释放 GIL，则可能是multiprocessing.dummy / concurrent.futures.ThreadPoolExecutor ，允许真正的线程并行性，由于 GIL，CPython 上的 CPU 绑定线程通常不可用）。