Numpy：将 3D 数组与 2D 数组相乘答案

【问题标题】：Numpy: multiply 3D array with 2D array [duplicate]Numpy：将 3D 数组与 2D 数组相乘
【发布时间】：2020-04-04 23:15:53
【问题描述】：

我有 2 个数组，一个形状为 (4, 2, 2)，另一个形状为 (4, 2).

A = np.random.rand(4, 2, 2)
B = np.random.rand(4, 2)

所以，对于每个0<= i <= 3，我都想拥有np.dot(A[i, :, :], B[i, :])。

如何在不使用 for 循环的情况下做到这一点？

【问题讨论】：

看这个：stackoverflow.com/questions/50796440/…

标签： python numpy

【解决方案1】：

(A @ B[...,np.newaxis])[...,0] 应该可以工作（np.matmul(A, B[...,np.newaxis])[...,0] 的语法糖）。

每the documentationmatmul：

如果任一参数为 N-D，N > 2，则将其视为位于最后两个索引中的矩阵堆栈。

你也可以使用 Einstein summation np.einsum("ijk,ik->ij", A, B)，虽然在我的机器上这显然比广播 matmul 慢：

import timeit

import numpy as np

A = np.random.rand(4, 2, 2)
B = np.random.rand(4, 2)


def broadcast_matmul():
    return (A @ B[..., np.newaxis])[..., 0]


def list_map():
    return np.array(list(map(np.dot, A, B)))


def einsum():
    return np.einsum("ijk,ik->ij", A, B)


print(timeit.timeit(broadcast_matmul))
print(timeit.timeit(list_map))
print(timeit.timeit(einsum))

# Results on my machine:
# 1.870921753
# 6.705698656999999
# 2.765732645

【讨论】：

【解决方案2】：

使用np.einsum 可以做到这一点：

np.einsum('ijk,ik->ij', A, B)

array([[0.39437083, 0.45360039],
       [0.75211742, 0.40669922],
       [0.18131254, 0.19473085],
       [0.2919673 , 0.10398859]])

使用相乘的第一个元素快速检查：

np.dot(A[0], B[0])
# array([0.39437083, 0.45360039])

【讨论】：

【解决方案3】：

这应该可以解决问题： list(map(np.dot, A, B))

【讨论】：