У меня есть два трехмерных тензора, тензор A, который имеет форму [B,N,S], и тензор B, который также имеет форму [B,N,S]. Я хочу получить третий тензор C, который, как я ожидаю, будет иметь форму [B,B,N], где находится элемент C[i,j,k] = np.dot(A[i,k,:], B[j,k,:]. Я тоже хочу добиться этого векторизованным способом.

Некоторая дополнительная информация: два тензора A и B имеют форму [Batch_size, Num_vectors, Vector_size]. Тензор C должен представлять собой скалярное произведение между каждым элементом в пакете из A и каждым элементом в пакете из B, между всеми различными векторами.

Надеюсь, что это достаточно ясно и с нетерпением жду ваших ответов!

4
gorjan 26 Июн 2019 в 16:32

3 ответа

Лучший ответ
In [331]: A=np.random.rand(100,200,300)                                                              
In [332]: B=A

Предлагаемый einsum, работающий непосредственно из

C[i,j,k] = np.dot(A[i,k,:], B[j,k,:] 

Выражение:

In [333]: np.einsum( 'ikm, jkm-> ijk', A, B).shape                                                   
Out[333]: (100, 100, 200)
In [334]: timeit np.einsum( 'ikm, jkm-> ijk', A, B).shape                                            
800 ms ± 25.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

matmul выполняет dot в последних 2 измерениях и рассматривает одно из первых как пакетное. В вашем случае «k» - это размерность пакета, а «m» - то, которое должно подчиняться правилу last A and 2nd to the last of B. Итак, переписываем ikm,jkm..., чтобы соответствовать, и транспонируем A и B соответственно:

In [335]: np.einsum('kim,kmj->kij', A.transpose(1,0,2), B.transpose(1,2,0)).shape                     
Out[335]: (200, 100, 100)
In [336]: timeit np.einsum('kim,kmj->kij',A.transpose(1,0,2), B.transpose(1,2,0)).shape              
774 ms ± 22.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Не большая разница в производительности. Но теперь используйте matmul:

In [337]: (A.transpose(1,0,2)@B.transpose(1,2,0)).transpose(1,2,0).shape                             
Out[337]: (100, 100, 200)
In [338]: timeit (A.transpose(1,0,2)@B.transpose(1,2,0)).transpose(1,2,0).shape                      
64.4 ms ± 1.17 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

И убедитесь, что значения совпадают (хотя чаще всего, если формы совпадают, значения соответствуют).

In [339]: np.allclose((A.transpose(1,0,2)@B.transpose(1,2,0)).transpose(1,2,0),np.einsum( 'ikm, jkm->
     ...:  ijk', A, B))                                                                              
Out[339]: True

Я не буду пытаться измерить использование памяти, но улучшение времени говорит о том, что оно тоже лучше.

В некоторых случаях einsum оптимизирован для использования matmul. Здесь, похоже, это не так, хотя мы могли бы поиграть с его параметрами. Я немного удивлен, что matmul работает намного лучше.

===

Я смутно припоминаю еще один SO о matmul сокращении, когда два массива - одно и то же, A@A. Я использовал B=A в этих тестах.

In [350]: timeit (A.transpose(1,0,2)@B.transpose(1,2,0)).transpose(1,2,0).shape                      
60.6 ms ± 1.17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [352]: B2=np.random.rand(100,200,300)                                                             
In [353]: timeit (A.transpose(1,0,2)@B2.transpose(1,2,0)).transpose(1,2,0).shape                     
97.4 ms ± 164 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Но это имело только скромное значение.

In [356]: np.__version__                                                                             
Out[356]: '1.16.4'

Мой BLAS и т. Д. - стандартный Linux, ничего особенного.

4
hpaulj 27 Июн 2019 в 15:49

Я думаю, что вы можете использовать einsum такие как:

np.einsum( 'ikm, jkm-> ijk', A, B)

С помощью индексов 'ikm, jkm-> ijk' вы можете указать, какое измерение будет сокращено в соответствии с соглашением Эйнштейна. Третье измерение обоих массивов A и B, называемое здесь 'm', будет уменьшено, так как операция dot выполняется для векторов.

2
Ben.T 26 Июн 2019 в 14:06

Пытаться:

C = np.diagonal( np.tensordot(A,B, axes=(2,2)), axis1=1, axis2=3)

От https://docs.scipy.org /doc/numpy/reference/generated/numpy.tensordot.html#numpy.tensordot

Объяснение

Решение состоит из двух операций. Сначала тензорное произведение между А и В над их третьей осью, как вы хотите. Это выдает тензор ранга 4, который вы хотите уменьшить до тензора ранга 3, взяв равные индексы по осям 1 и 3 (ваш k в вашей нотации, обратите внимание, что tensordot дает другую ось порядок, чем ваша математика). Это можно сделать, взяв диагональ, как это можно сделать при уменьшении матрицы до вектора ее диагональных элементов.

-1
Learning is a mess 26 Июн 2019 в 14:00