张量流中两个相同等级张量之间的广播答案

【问题标题】：Broadcasting between two same-rank tensors in tensorflow张量流中两个相同等级张量之间的广播
【发布时间】：2017-09-01 02:48:08
【问题描述】：

我有两个张量 x 和 s 的形状：

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

我想通过维度1广播x和s之间的点积如下：

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

在哪里

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)])

我有这个实现：

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]

我知道这在内存中效率不高，因为 tile。此外，reshape's、transpose's element-wise 和 reduce_sums 操作可能会损害较大张量的性能。有没有其他方法可以让它更干净？

【问题讨论】：

标签： tensorflow array-broadcasting tensorflow-xla

【解决方案1】：

你有任何证据表明reshapes 很贵吗？以下使用重塑和维度广播：

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

【讨论】：

感谢您的帮助。实际上，您的 impl (x_s2) 比 OP (x_s) > %timeit sess.run(x_s) 1000 个循环快，最好的 3：每个循环 365 µs > %timeit sess.run(x_s2) 1000 个循环，最好的 3：每个循环 243 µs

【解决方案2】：

只是一些建议，也许不会比你的更快。先将s和tf.split拆分成四个张量，然后用tf.tensordot得到最终结果，像这样

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)

【讨论】：