使用列表索引时如何加速“for”循环？（Python）

Question

我尝试使用 Numpy 函数或向量而不是 for 循环来加速此代码：sommes = []for j in range(vertices.shape[0]):    terme = new_vertices[j] - new_vertices[vertex_neighbors[j]]    somme_j = np.sum(terme)    sommes.append(somme_j)E_int = np.sum(sommes)（它是迭代算法的一部分，并且有很多“顶点”，所以我认为 for 循环花费的时间太长。）例如，要计算 j = 0 时的“terme”，我有：In: new_vertices[0]Out: array([ 10.2533888 , -42.32279717,  68.27230793])In: vertex_neighbors[0]Out: [1280, 2, 1511, 511, 1727, 1887, 759, 509, 1023]In: new_verties[vertex_neighbors[0]]Out: array([[ 10.47121043, -42.00123956,  68.218715  ],            [ 10.2533888 , -43.26905874,  62.59473849],            [ 10.69773735, -41.26464083,  68.09594854],            [ 10.37030712, -42.16729601,  68.24639107],            [ 10.12158146, -42.46624547,  68.29621598],            [  9.81850836, -42.71158695,  68.33710623],            [  9.97615447, -42.59625943,  68.31788497],            [ 10.37030712, -43.11676015,  62.54960623],            [ 10.55512696, -41.82622703,  68.18954624]])In: new_vertices[0] - new_vertices[vertex_neighbors[0]]Out: array([[-0.21782162, -0.32155761,  0.05359293],             [ 0.        ,  0.94626157,  5.67756944],             [-0.44434855, -1.05815634,  0.17635939],             [-0.11691832, -0.15550116,  0.02591686],             [ 0.13180734,  0.1434483 , -0.02390805],             [ 0.43488044,  0.38878979, -0.0647983 ],             [ 0.27723434,  0.27346227, -0.04557704],             [-0.11691832,  0.79396298,  5.7227017 ],             [-0.30173816, -0.49657014,  0.08276169]])问题是 new_vertices[vertex_neighbors[j]] 并不总是具有相同的大小。例如，当 j = 7 时：In: new_vertices[7]Out: array([ 10.74106112, -63.88592276, -70.15593947])In: vertex_neighbors[7]Out: [1546, 655, 306, 1879, 920, 925]没有for循环可以吗？我的想法已经用完了，所以任何帮助将不胜感激！

Answer 1

是的，这是可能的。这个想法是用来np.repeat创建一个向量，其中的项目重复可变次数。这是代码：

# The two following lines can be done only once if the indices are constant between iterations (precomputation)

counts = np.array([len(e) for e in vertex_neighbors])

flatten_indices = np.concatenate(vertex_neighbors)

E_int = np.sum(np.repeat(new_vertices, counts, axis=0) - new_vertices[flatten_indices])

这是一个基准：

import numpy as np

from time import *

n = 32768

vertices = np.random.rand(n, 3)

indices = []

count = np.random.randint(1, 10, size=n)

for i in range(n):

    indices.append(np.random.randint(0, n, size=count[i]))

def initial_version(vertices, vertex_neighbors):

    sommes = []

    for j in range(vertices.shape[0]):

        terme = vertices[j] - vertices[vertex_neighbors[j]]

        somme_j = np.sum(terme)

        sommes.append(somme_j)

    return np.sum(sommes)

def optimized_version(vertices, vertex_neighbors):

    # The two following lines can be precomputed

    counts = np.array([len(e) for e in indices])

    flatten_indices = np.concatenate(indices)

    return np.sum(np.repeat(vertices, counts, axis=0) - vertices[flatten_indices])

def more_optimized_version(vertices, vertex_neighbors, counts, flatten_indices):

    return np.sum(np.repeat(vertices, counts, axis=0) - vertices[flatten_indices])

timesteps = 20

a = time()

for t in range(timesteps):

    res = initial_version(vertices, indices)

b = time()

print("V1: time:", b - a)

print("V1: result", res)

a = time()

for t in range(timesteps):

    res = optimized_version(vertices, indices)

b = time()

print("V2: time:", b - a)

print("V2: result", res)

a = time()

counts = np.array([len(e) for e in indices])

flatten_indices = np.concatenate(indices)

for t in range(timesteps):

    res = more_optimized_version(vertices, indices, counts, flatten_indices)

b = time()

print("V3: time:", b - a)

print("V3: result", res)

这是我机器上的基准测试结果：

V1: time: 3.656714916229248

V1: result -395.8416223057596

V2: time: 0.19800186157226562

V2: result -395.8416223057595

V3: time: 0.07983255386352539

V3: result -395.8416223057595

正如您所看到的，这一优化版本比参考实现快 18 倍，而预先计算索引的版本比参考实现快 46 倍。

请注意，优化版本应该需要更多 RAM（特别是如果每个顶点的邻居数量很大）。