Broadcasting entre deux tenseurs de même rang dans tensorflow

Question

J'ai deux tenseurs x et s avec des formes:

> x.shape 
TensorShape([Dimension(None), Dimension(3), Dimension(5), Dimension(5)]) 
> s.shape 
TensorShape([Dimension(None), Dimension(12), Dimension(5), Dimension(5)]) 

Je veux diffuser le produit scalaire entre x et s à travers la dimension 1 comme suit:

> x_s.shape 
TensorShape([Dimension(None), Dimension(4), Dimension(5), Dimension(5)]) 

où

x_s[i, 0, k, l] = sum([x[i, j, k, l] * s[i, j, k, l] for j in range (3)]) 
x_s[i, 1, k, l] = sum([x[i, j-3, k, l] * s[i, j, k, l] for j in range (3, 6)]) 
x_s[i, 2, k, l] = sum([x[i, j-6, k, l] * s[i, j, k, l] for j in range (6, 9)]) 
x_s[i, 3, k, l] = sum([x[i, j-9, k, l] * s[i, j, k, l] for j in range (9, 12)]) 

J'ai cette implémentation:

s_t = tf.transpose(s, [0, 2, 3, 1]) # [None, 5, 5, 12] 
x_t = tf.transpose(x, [0, 2, 3, 1]) # [None, 5, 5, 3] 
x_t = tf.tile(x_t, [1, 1, 1, 4]) # [None, 5, 5, 12] 

x_s = x_t * s_t # [None, 5, 5, 12] 
x_s = tf.reshape(x_s, [tf.shape(x_s)[0], 5, 5, 4, 3]) # [None, 5, 5, 4, 3] 
x_s = tf.reduce_sum(x_s, axis=-1) # [None, 5, 5, 4] 
x_s = tf.transpose(x_s, [0, 3, 1, 2]) # [None, 4, 5, 5] 

Je comprends que ce n'est pas efficace en mémoire à cause de la tile. En outre, reshape, transposeelement-wise et reduce_sum s opérations peuvent nuire à la performance pour les tenseurs plus grands. Y a-t-il une alternative pour le rendre plus propre?

Answer 1

Avez-vous des preuves que reshape s sont chers? Ce qui suit utilise une diffusion remodeler et de dimension:

x_s = tf.reduce_sum(tf.reshape(s, (-1, 4, 3, 5, 5)) * 
        tf.expand_dims(x, axis=1), axis=2) 

Answer 2

Juste quelques conseils et peut-être pas plus rapide que le vôtre. Tout d'abord divisé s avec tf.split en quatre tenseurs, puis utilisez tf.tensordot pour obtenir résultat final, comme ce

splits = tf.split(s, [3] * 4, axis=1) 
splits = map(lambda split: tf.tensordot(split, x, axes=[[1], [1]]), splits) 
x_s = tf.stack(splits, axis=1)