python - 轴上的 Numpy 乘法

Question

我有 2 个 numpy 数组，一个是 shape (2, 5, 10)，另一个是 shape (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.

标签： pythonnumpy

Answer 1

这样做的一种方法，假设您对xas type很好np.array，是通过更好地利用NumPy 广播（参见@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)