Question

我尝试运行这个小代码，它只需要随机点（这里是 50k，接近我实际拥有的点）并返回随机选择的每个点的第 10 个最近点。

但不幸的是，这（真的！）很长，因为肯定是循环。

由于我对“代码优化”还很陌生，有什么技巧可以让这一切变得更快？（在 Python 规模上更快，我知道我不是用 C++ 编码）。

这是一个可重现的示例，其数据大小接近我所拥有的数据：

import time

import numpy as np
from numpy import random
from scipy.spatial import distance

# USEFUL FUNCTION

start_time = time.time()


def closest_node(node, nodes):
    nodes = np.asarray(nodes)
    deltas = nodes - node
    dist_2 = np.einsum("ij,ij->i", deltas, deltas)
    ndx = dist_2.argsort()
    return data[ndx[:10]]


# REPRODUCIBLE DATA

mean = np.array([0.0, 0.0, 0.0])
cov = np.array([[1.0, -0.5, 0.8], [-0.5, 1.1, 0.0], [0.8, 0.0, 1.0]])
data = np.random.multivariate_normal(mean, cov, 500000)

# START RUNNING

points = data[np.random.choice(data.shape[0], int(np.round(0.1 * len(data), 0)))]
print(len(points))

for w in points:
    closest_node(w, data)

print("--- %s seconds ---" % (time.time() - start_time))

Answer 1

在 500000 个元素的数组上运行 argsort 的每个循环所需的时间是巨大的。我能想到的唯一改进是使用可以返回最小的 10 个元素而无需对整个数组进行完全排序的东西。

A fast way to find the largest N elements in an numpy array

所以代替

ndx = dist_2.argsort() 
return data[ndx[:10]]

应该是

ndx = np.argpartition(dist_2, 10)[:10]
return data[ndx[:10]]

我只对 500 点进行了基准测试，因为在我的 PC 上运行已经花费了相当长的时间。

N=500
Using argsort: 25.625439167022705 seconds
Using argpartition: 6.637120485305786 seconds

Answer 2

您最好通过分析器分析最慢的点：How do I find out what parts of my code are inefficient in Python

乍一看可能的一件事是，您应该尽可能多地移动到循环外。如果您打算通过 np.asarray() 转换点，最好在循环之前对所有点执行一次，并在函数中使用结果，而不是在每次循环运行中执行 np.asarray() 。

Python 循环效率低下，尽管它很简单

2 个答案: