張量流中的廣播點積

Question

在張量流中，我有以下問題。

我有一個張量米形狀的 [batch_size時，dim_a，dim_b]和基質ü形狀[batch_size時，dim_b]的。

M = tf.constant(shape=[batch_size, sequence_size, embed_dim]) 
U = tf.constant(shape=[batch_size, embed_dim]) 

我該怎麼實現的是[我，dim_a，dim_b]×[我，dim_b]我批各指標的點積。

P[i] = tf.matmul(M[i, :, :], tf.expand_dims(U[i, :], 1)) for each i. 

基本上，broadacast在批處理軸上的點積。這是可能的，我該如何實現它？

Answer 1

這可以用tf.einsum（）來實現：

import tensorflow as tf 
import numpy as np 

batch_size = 2 
sequence_size = 3 
embed_dim = 4 

M = tf.constant(range(batch_size * sequence_size * embed_dim), shape=[batch_size, sequence_size, embed_dim]) 
U = tf.constant(range(batch_size, embed_dim), shape=[batch_size, embed_dim]) 

prod = tf.einsum('bse,be->bs', M, U) 

with tf.Session(): 
    print "M" 
    print M.eval() 
    print 
    print "U" 
    print U.eval() 
    print 
    print "einsum result" 
    print prod.eval() 
    print 

    print "numpy, example 0" 
    print np.matmul(M.eval()[0], U.eval()[0]) 
    print 
    print "numpy, example 1" 
    print np.matmul(M.eval()[1], U.eval()[1]) 

輸出：

M 
[[[ 0 1 2 3] 
    [ 4 5 6 7] 
    [ 8 9 10 11]] 

[[12 13 14 15] 
    [16 17 18 19] 
    [20 21 22 23]]] 

U 
[[2 3 3 3] 
[3 3 3 3]] 

einsum result 
[[ 18 62 106] 
[162 210 258]] 

numpy, example 0 
[ 18 62 106] 

numpy, example 1 
[162 210 258]