一种使用 numpy 计算批处理和时间维度空间变换的更快方法

Question

对于机器学习任务，我必须为时间序列中的每一帧计算 3D 坐标的空间转换(在 z 轴上的旋转)。而且，我有一批这样的时间序列。我想尽可能减少 for 循环的使用。

假设我有一个形状的旋转矩阵(batch_size，3,3)和一个形状的张量(batch_size，seq_length，n_coordinates，3)。我目前正在做的是随时间和批量尺寸进行双倍循环，并计算每个 3D 坐标的点积。

这是代码：

# Compute transformation
for t in range(seq_length):
for b in range(batch_size):
X[b, t, :, :] = np.dot(rotation_z_matrix[b], X[b, t, :, :].T).T

我已经看过张量点和 einsum 函数，但最终，我不想在一个维度上对点积求和，我想在 2 维(批次和时间)上堆叠我的点积。

是否有有效的等待来编写等效代码？

提前谢谢！

Answer 1

您可以使用np.einsum-

X_new = np.einsum('ijk,ilmk->ilmj',rotation_z_matrix,X)

此外，在np.einsum中optimize标志中，将其设置为使用 BLAS 的True。

Answer 2

可以使用广播来完成：

X@rotation_z_matrix.transpose(0,2,1)[:, None, ...]

这给出了(在模拟数据集上)与@Divakar相同的答案

batch_size = 10
seq_length = 8
n_coordinates = 12
X = np.random.randint(0,10,(batch_size, seq_length, n_coordinates, 3))
rotation_z_matrix = np.random.randint(0,10,(batch_size,3,3))
(X@rotation_z_matrix.transpose(0,2,1)[:, None, ...] == np.einsum('ijk,ilmk->ilmj',rotation_z_matrix,X)).all()
# True

但至少对于此示例，它要快得多。

timeit(lambda: np.einsum('ijk,ilmk->ilmj',rotation_z_matrix,X, optimize=True), number=1000)
# 0.1285447319969535
timeit(lambda: np.einsum('ijk,ilmk->ilmj',rotation_z_matrix,X, optimize=False), number=1000)
# 0.07962286799738649
timeit(lambda: X@rotation_z_matrix.transpose(0,2,1)[:, None, ...], number=1000)
# 0.019039910010178573

请务必注意，设置optimize标志实际上会减慢einsum速度。(这种情况经常发生在我身上。

更新：相同的示例，但数据转换为浮点型 dtype

timeit(lambda: np.einsum('ijk,ilmk->ilmj',rotation_z_matrix,X, optimize=True), number=1000)
# 0.12346570500812959
timeit(lambda: np.einsum('ijk,ilmk->ilmj',rotation_z_matrix,X, optimize=False), number=1000)
# 0.07575376800377853
timeit(lambda: X@rotation_z_matrix.transpose(0,2,1)[:, None, ...], number=1000)
# 0.027829282989841886