跨轴的Keras / Tensorflow批量矩阵乘法

时间:2018-06-11 13:12:50

标签: tensorflow matrix keras deep-learning

说我有张量

a
Out[15]: <tf.Tensor 'Placeholder_2:0' shape=(?, 1152, 8) dtype=float32>
b
Out[16]: <tf.Variable 'Variable:0' shape=(16, 8, 1152, 10) dtype=float32_ref>

表示一批1152个8维向量和 b是1152 * 10,(16,8)矩阵。

我希望将这些矩阵与a中的8维向量相乘,得到一个形状的张量(None,16,1152,10)。我知道在张量流中,可以使用einsum来完成这项任务

tf.einsum('ijkl,bkj->bikl', b, a)

给我正确的输出和形状。但与tf.einsumK.batch_dot等类似功能相比,tf.tensordot速度非常慢。但是,我很难理解这些函数如何处理轴和广播规则。有什么帮助吗?

1 个答案:

答案 0 :(得分:1)

通过使用transposereshape,您可以实现相同目标:

a : [batch, 1152, 8] --> reshape --> [batch, 1, 1, 1152, 8]
b : [16,8,1152,10]   --> transpose --> [16, 10, 1152, 8] 
                     --> expand_dims --> [1, 16, 10, 1152, 8]
multiply (a, b)      --> [batch, 16, 10, 1152, 8] 
reduce_sum axis 4    --> [batch, 16, 10, 1152]             

代码:

#inputs
import numpy.testing as npt
x = np.random.normal(size=(5,1152,8)) 
y = np.random.normal(size=(16, 8, 1152, 10))

a = tf.placeholder(tf.float32,shape=(None, 1152, 8))
b = tf.constant(y, tf.float32)

out = tf.reduce_sum(tf.expand_dims(tf.transpose(b,[0, 3, 2, 1]),0) 
                   * tf.reshape(a,[-1,1,1,tf.shape(a)[1], tf.shape(a)[2]]), axis=4)
out = tf.transpose(out, [0,1,3,2])
out_ein = tf.einsum('ijkl,bkj->bikl', b, a)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    o = sess.run(out, {a: x})
    e = sess.run(out_ein, {a: x})
    npt.assert_almost_equal(o, e, decimal=5)
    #almost the same