Einsum медленнее, чем явная реализация Numpy для тензорно-матричного произведения n-мод

Question

Я пытаюсь реализовать тензорно-матричный продукт n-режима (как определено Kolda и Bader: https://www.sandia.gov/~tgkolda/pubs/pubfiles/SAND2007-6702.pdf) эффективно в Python с использованием Numpy. Операция эффективно сводится к (для матрицы U, тензор Xи ось / режим k):

1. Извлечение всех векторов вдоль оси k из X путем свертывания всех других осей.

2. Умножьте эти векторы наслева от U с использованием стандартного умножения матриц.

3. Вставьте векторы снова в выходной тензор, используя ту же форму, за исключением X.shape [k], который теперь равен U.shape[0] (изначально X.shape [k] должен быть равен U.shape [1] в результате умножения матрицы).

Я использовалявная реализация на некоторое время, которая выполняет все эти шаги отдельно:

1. Транспонировать тензор, чтобы вывести ось k вперед (в моем полном коде я добавил исключение в случае k == X.ndim - 1, в этом случае быстрее оставить его там и перенести всю будущую работуили, по крайней мере, в моем приложении, но здесь это не имеет значения).

2. Измените тензор, чтобы свернуть все остальные оси.

3. Рассчитатьумножение матриц.

4. Измените тензор, чтобы восстановить все остальные оси.

5. Переместите тензор обратно в исходный порядок.

Я думаю, что эта реализация создает множество ненужных (больших) массивов, поэтому, обнаружив np.einsum, я подумал, что это значительно ускорит процесс.Однако, используя приведенный ниже код, я получил худшие результаты:

import numpy as np
from time import time

def mode_k_product(U, X, mode):
    transposition_order = list(range(X.ndim))
    transposition_order[mode] = 0
    transposition_order[0] = mode
    Y = np.transpose(X, transposition_order)
    transposed_ranks = list(Y.shape)
    Y = np.reshape(Y, (Y.shape[0], -1))
    Y = U @ Y
    transposed_ranks[0] = Y.shape[0]
    Y = np.reshape(Y, transposed_ranks)
    Y = np.transpose(Y, transposition_order)
    return Y

def einsum_product(U, X, mode):
    axes1 = list(range(X.ndim))
    axes1[mode] = X.ndim + 1
    axes2 = list(range(X.ndim))
    axes2[mode] = X.ndim
    return np.einsum(U, [X.ndim, X.ndim + 1], X, axes1, axes2, optimize=True)

def test_correctness():
    A = np.random.rand(3, 4, 5)
    for i in range(3):
        B = np.random.rand(6, A.shape[i])
        X = mode_k_product(B, A, i)
        Y = einsum_product(B, A, i)
        print(np.allclose(X, Y))

def test_time(method, amount):
    U = np.random.rand(256, 512)
    X = np.random.rand(512, 512, 256)
    start = time()
    for i in range(amount):
        method(U, X, 1)
    return (time() - start)/amount

def test_times():
    print("Explicit:", test_time(mode_k_product, 10))
    print("Einsum:", test_time(einsum_product, 10))

test_correctness()
test_times()

Время для меня:

Явное: 3.9450525522232054

Einsum: 15.873924326896667

Это нормально?или я что то не так делаю?Я знаю, что существуют обстоятельства, когда сохранение промежуточных результатов может снизить сложность (например, умножение цепочки матриц), однако в этом случае я не могу думать о каких-либо вычислениях, которые повторяются.Является ли матричное умножение настолько оптимизированным, что оно устраняет преимущества отсутствия транспонирования (что технически имеет меньшую сложность)?

Answer 1

Я более знаком со стилем подписок при использовании einsum, поэтому сработал следующие эквивалентности:

In [194]: np.allclose(np.einsum('ij,jkl->ikl',B0,A), einsum_product(B0,A,0))          
Out[194]: True
In [195]: np.allclose(np.einsum('ij,kjl->kil',B1,A), einsum_product(B1,A,1))          
Out[195]: True
In [196]: np.allclose(np.einsum('ij,klj->kli',B2,A), einsum_product(B2,A,2))          
Out[196]: True

С параметром mode ваш подход в einsum_product может быть лучшим.Но эквивалентности помогают мне лучше визуализировать вычисления и могут помочь другим.

Время должно быть в основном таким же.В einsum_product есть дополнительное время настройки, которое должно исчезнуть в больших измерениях.

Answer 2

После обновления Numpy Einsum лишь немного медленнее, чем явный метод, с многопоточностью или без нее (см. Комментарии к моему вопросу).