主页 / coding / 问题 / 79173190
Accepted
Henkie
Asked: 2024-11-10 00:20:19 +0800 CST

这段 Python 代码可以加速吗？

  • 772

我有以下想要加速的代码：

import numpy as np

def do_calc(a, b, d, sum_)->None:
    idx = np.ndindex(a.shape[:1]) 
    sum_[:] = 0
    for i in idx:
       sum_[0:d[i[0]]] -= a[i, 0:d[i[0]]][0]
       sum_[d[i[0]]] += b[i, d[i[0]]][0]
    return sum_


a=np.ones((5,10), order='F')
b=np.ones((5,10), order='F')*10

sum_ = np.empty(a.shape[1])

d = [8,5,2,7,1]
do_calc(a, b, d, sum_)

因此 a 和 b 是静态的 5x10 数组，我想将它们加起来形成sum_一个长度为 1d 的数组，如下所示：

  • 对于每一行i，b我只想包含索引d[i]，也就是b[i, d[i]]将进入位置sum_[d[i]]

  • 对于每一行i，a我想包含所有到索引的值d[i]，也就是说a[i, 0:d[i]]，将进入位置sum_[0:d[i]]

在这个例子中，我们有：


a= array([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

b= array([[10., 10., 10., 10., 10., 10., 10., 10., 10., 10.],
          [10., 10., 10., 10., 10., 10., 10., 10., 10., 10.],
          [10., 10., 10., 10., 10., 10., 10., 10., 10., 10.],
          [10., 10., 10., 10., 10., 10., 10., 10., 10., 10.],
          [10., 10., 10., 10., 10., 10., 10., 10., 10., 10.]])

d= [8, 5, 2, 7, 1]

sum_= array([-5.,  6.,  7., -3., -3.,  8., -2.,  9., 10.,  0.])

我尝试了以下方法：

  • 使用 Fortan 顺序，因为我们主要复制最后一个索引
  • sum_传递给 do_calc 函数，这样就不必每次都启动它。最后两行有一个大循环，我们每次都会为 d 生成新值并运行 do_calc。
  • 使用 ndindex 来迭代行。

我想知道是否可以更有效地完成此操作，因为代码一直在处理这一部分。而实际上 a 和 b 不是常量，因此我们无法在此示例中利用这一点。

python

1 个回答

  1. Best Answer

    如果您能提供有关此类功能需求的背景信息，我很乐意调整我的回复，但考虑到您的功能，我能够将代码简化为以下内容。

    def docal2(a: np.array, b: np.array, d: np.array, sum_: np.array) -> np.array:
    r, c = a.shape
    sum_[:] = 0
    for i in range(r):
        ii = d[i]
        sum_[0:ii] -= a[i, 0: ii]
        sum_[ii] += b[i, ii]
    return sum_

    这应该可以防止对原始数据进行不必要的切片，并显著加快计算速度。如果您可以提供进一步的背景信息，我很乐意帮助您对计算进行矢量化，而不是原始迭代。

    几点：

    • a[i, 0:d[i[0]]][0]相当于a[i, 0:i]，因为你总是需要第 0 个元素来进行计算。同样适用于b[i, d[i[0]]][0]
    • 而不是np.ndindex(a.shape[:1])使用range来简化下游索引i[0]。您可以直接使用i进行索引。

    与原函数比较

    def timeit(func):
    import time
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        return result, time.time() - start
    return wrapper

    测试：

    loss_holder = dict()


for rng in range(100, 10_000, 100):
    sub_holder = []
    for _ in range(10):
        a = np.ones((rng, 10)) * 1.
        b = np.ones((rng, 10)) * 10.
        d = np.random.choice(10, rng)
        
        sum_= np.array([-5.,  6.,  7., -3., -3.,  8., -2.,  9., 10.,  0.])
        
        r1, t1 = do_calc(a, b, d, sum_)
        r2, t2 = docal2(a, b, d, sum_)

        if not all(r1 == r2):
            raise ValueError(f"r1: {r1}, r2: {r2}")
        
        loss = t1 / (t2 if t2 != 0 else 1e-100)
        sub_holder.append(loss)

    loss_holder[rng] = np.mean(sub_holder)


import matplotlib.pyplot as plt


plt.plot(loss_holder.keys(), loss_holder.values())
plt.ylim(0, 6)

    在此处输入图片描述

    • 1

相关问题