Theano: Как взять «матричное внешнее произведение», где элементы являются матрицами

Question

В принципе, у меня есть 2 тензора: A, где A.shape = (N, H, D) и B, где B.shape = (K, H, D). То, что я хотел бы сделать, это получить тензор, С, с формой (N, K, D, H) таким образом, что:

C[i, j, :, :] = A[i, :, :] * B[j, :, :]. 

Может ли это быть сделано эффективно Теано?

Side Примечание: Фактический конечный результат, который я хотел бы добиться того, чтобы иметь тензор, Е, формы (N, K, D), что:

E[i, j, :] = (A[i, :, :]*B[j, :, :]).sum(0) 

Итак, если есть способ, чтобы получить это прямо, Я бы предпочел это (с надеждой спасает место).

Answer 1

Один подхода можно предположить, что использует broadcasting -

(A[:,None]*B).sum(2) 

Пожалуйста, обратите внимание, что создается будет иметь форму (N, K, H, D) перед тем суммарного сокращения на axis=2 промежуточного массив снижает его (N,K,D).

Answer 2

Вы можете получить окончательный трехмерный результат E без создания большой промежуточной матрицы с использованием batched_dot:

import theano.tensor as tt 
A = tt.tensor3('A') # A.shape = (D, N, H) 
B = tt.tensor3('B') # B.shape = (D, H, K) 
E = tt.batched_dot(A, B) # E.shape = (D, N, K) 

К сожалению, для этого вам необходимо переставить размеры на входные и выходные массивы. Хотя это может быть сделано с dimshuffle в Теано, кажется batched_dot не может справиться с произвольно strided массивами и поэтому следующее поднимает ValueError: Some matrix has no unit stride когда E оценивается:

import theano.tensor as tt 
A = tt.tensor3('A') # A.shape = (N, H, D) 
B = tt.tensor3('B') # B.shape = (K, H, D) 
A_perm = A.dimshuffle((2, 0, 1)) # A_perm.shape = (D, N, H) 
B_perm = B.dimshuffle((2, 1, 0)) # B_perm.shape = (D, H, K) 
E_perm = tt.batched_dot(A_perm, B_perm) # E_perm.shape = (D, N, K) 
E = E_perm.dimshuffle((1, 2, 0)) # E.shape = (N, K, D) 

batched_dot использует scan вдоль первой (размер D) измерение , Поскольку scan выполняется последовательно, это может быть вычислительно менее эффективным, чем вычисление всех продуктов параллельно при работе на графическом процессоре.

Вы можете обменять между эффективностью использования памяти в соответствии с методом приближения и параллелизма в широковещательном подходе с использованием scan. Идея в том, чтобы вычислить полный продукт C для партий размером M параллельно (предполагая, что M является точным фактором D), итерация партиями с scan:

import theano as th 
import theano.tensor as tt 
A = tt.tensor3('A') # A.shape = (N, H, D) 
B = tt.tensor3('B') # B.shape = (K, H, D) 
A_batched = A.reshape((N, H, M, D/M)) 
B_batched = B.reshape((K, H, M, D/M)) 
E_batched, _ = th.scan(
    lambda a, b: (a[:, :, None, :] * b[:, :, :, None]).sum(1), 
    sequences=[A_batched.T, B_batched.T] 
) 
E = E_batched.reshape((D, K, N)).T # E.shape = (N, K, D)