如何使用numpy优化循环参数化函数调用?

Question

我正在尝试将numpy数组上的嵌套循环转换为numpy优化的实现。在循环内部调用的函数接受一个4D向量和一个单独的参数，并输出一个4D向量，该向量应该根据新值的操作替换旧的4D向量。如果相关，则函数为Welford在线更新，根据新值更新均值和标准差，4D向量为[old_mean, old_std, old_s, num_values]。对于每个像素通道，我将这些值保存在history_array中，以便根据未来的像素值更新分布。

我现在的代码是这样的:

def welford_next(arr:np.ndarray, new_point:np.float32) -> np.ndarray:
old_mean, _, old_s, num_points = arr
num_points += 1
new_mean = old_mean + (new_point - old_mean) / num_points
new_s = old_s + (new_point - old_mean) * (new_point - new_mean)
return [new_mean, np.sqrt(new_s / num_points) if num_points > 1 else new_s, new_s, num_points]
updates = [10., 20., 30., 40., 90., 80.]
history_array = np.zeros(shape = b.shape + (4,))   # shape: [6,3,3,4]
print(f'History Shape: {history_array.shape}')
history_array_2 = np.zeros_like(history_array)
for update in updates:
image = np.empty(shape = b.shape)              # shape: [6,3,3] (h x w x c)
image.fill(update)
for i, row in enumerate(image):                # Prohibitively expensive
for j, col in enumerate(row):
for k, channel in enumerate(col):
history_array[i][j][k] = welford_next(history_array[i][j][k], channel)
history_array_2 = np.apply_along_axis(welford_next, axis=2, arr=history_array_2)

print(history_array == history_array_2)

然而，np.apply_along_axis()似乎不可行，因为它不允许与数组本身一起传递额外的参数。我也遇到了np.ufunc,welford_next()功能可以转换为使用np.frompyfunc()，但目前还不清楚它如何帮助我达到预期的目标。

如何使用numpy实现这个循环操作?

Answer 1

numpy优化的方法是改变我们使用welford_next()函数的方式。正如评论中提到的，对函数的重复调用无法优化，因此函数调用需要限制为每帧一次，并且需要在函数本身内部进行优化。下面的实现工作速度要快50倍。

def welford(history:np.ndarray, frame:np.ndarray) -> np.ndarray:
old_mean, _, old_s, num_points = np.transpose(history, [3,0,1,2])
num_points += 1.
new_mean = old_mean + (frame - old_mean) / num_points
new_s = old_s + (frame - old_mean) * (frame - new_mean)
new_std = np.sqrt(new_s / num_points) if num_points[0][0][0] > 1 else new_s
return np.transpose(np.array([new_mean, new_std, new_s, num_points]), [1,2,3,0])
updates = [10., 20., 30., 40., 90., 80.]
history_array = np.zeros(shape = b.shape + (4,))   # shape: [6,3,3,4]
for update in updates:
image = np.empty(shape = b.shape)              # shape: [6,3,3] (h x w x c)
image.fill(update)
history_array = welford(history_array, image)