Более быстрый способ вычисления пространственного преобразования по пакетному и временному измерению с помощью NumPy

Question

Для задачи машинного обучения мне необходимо вычислить пространственное преобразование трехмерных координат (вращение по оси z) для каждого кадра во временной последовательности. Более того, у меня есть серия этих временных рядов. Я хотел бы максимально сократить использование циклов for.

Скажем, у меня есть матрица вращения формы (batch_size, 3, 3) и тензор формы (batch_size, seq_length, n_coordinates, 3). В настоящее время я делаю двойной цикл по времени и размерность пакета и вычисление точечного произведения для каждой трехмерной координаты.

Вот код:

# 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

Я рассмотрел функции тензордот и эйнсум, но в итоге я не хочу суммировать точечные продукты по измерению, я хочу сложить свои точечные продукты по 2 измерениям (партиям и времени).

Есть ли время ожидания, чтобы написать эквивалентный код?

Заранее спасибо!

Answer 1

Вы можете использовать np.einsum -

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

Кроме того, поиграйте с флагом optimize в np.einsum, установив его как True для использования BLAS.

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