Question

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

说我有一个形状为（batch_size，3，3）的旋转矩阵和形状为（batch_size，seq_length，n_coordinates，3）的张量。我目前正在做的是在时间和批次尺寸上进行两次for循环，并为每个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)

另外，通过将其设置为optimize以使用BLAS，在np.einsum中使用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的速度。（这种情况经常发生在我身上。）

更新：同一示例，但数据转换为float 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

使用numpy

2 个答案: