python - Numpy 轴上的乘法

Question

我有 2 个 numpy 数组，一个的形状为 (2, 5, 10)，另一个的形状为 (2, 5, 10, 10) 。我想要的乘法是[[x1[i][j] * x2[i][j] for j in range(5)] for i in range(2)]它按预期工作，但速度相当慢，我想直接乘以 x1 * x2 但 numpy 不喜欢这样。是否有 numpy 方法可以在给定轴上相乘？

我尝试过 numpy.multiply，因为它说“axis”对于 ufunc“multiply”来说是无效关键字

x1 = np.arange(100).reshape((2, 5, 10))
x2 = np.arange(1000).reshape((2, 5, 10, 10))
x = [[x1[i][j] * x2[i][j] for j in range(5)] for i in range(2)] # slow method that works
x = np.multiply(x1, x2, axis=2) # What I'm looking for but doesn't work.

Answer 1

假设您可以使用 x 作为类型 np.array，实现此目的的一种方法是更好地利用 NumPy broadcasting (参见@hpaulj的回答):

import numpy as np

x1 = np.arange((2 * 3 * 4)).reshape((2, 3, 4))
x2 = np.arange((2 * 3 * 4 * 4)).reshape((2, 3, 4, 4))
x = [[x1[i][j] * x2[i][j] for j in range(3)] for i in range(2)]) # slow method that works

# : using NumPy broadcasting
y = x1[:, :, None, :] * x2

np.all(np.array(x) == y)
# True

Timewise 加速了约 10 倍:

%timeit np.array([[x1[i][j] * x2[i][j] for j in range(3)] for i in range(2)])
# 13.6 µs ± 388 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit x1[:, :, None, :] * x2
# 1.4 µs ± 10.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)