Умножение матриц по оси теноров с numpy и с графическим процессором

Question

У меня есть матрица X с формой (F,T,M). Я умножу sh, чтобы умножить каждую матрицу (T,M) вдоль оси F, чтобы получить ответ в форме (M,M,F). Этот код выполняет свою работу для меня, но эта операция повторяется много раз и очень медленно:

    for f in range(F):
        output[:,:,f] = np.matmul(X[f,:,:].T,X[f,:,:])

Все, что я мог найти, это np.tensordot () функция. Если я правильно понимаю, это не очень хороший вариант для меня, так как мне нужно умножение матриц, а не скалярное произведение.

Как эффективно реализовать это, используя numpy? Возможно ли и выгодно ли использовать keras \ tf для этой цели?

Answer 1

Мы можем использовать np.matmul/@ opeartor in Python 3.x после расширения размеров -

np.matmul(X.swapaxes(1,2),X).swapaxes(0,2)
(X.swapaxes(1,2)@X).swapaxes(0,2)

В качестве альтернативы, с np.einsum с прямым переводом переменных формы, используемых для обозначения строки -

np.einsum('ftm,ftn->mnf',X,X)