Question

说我有两个张量：

a=Tensor("zeros_3:0", shape=(2, 4, 5), dtype=float32)
b=Tensor("ones_3:0", shape=(2, 3, 5), dtype=float32)

如何使用嵌套的map_fn或其他tf函数沿轴2合并每个元素以获得新的(2,3,4,10)形张量？

这是我的for循环版本

        concat_list = []
        for i in range(a.get_shape()[1]):
            for j in range(b.get_shape()[1]):
                concat_list.append(tf.concat([a[:, i, :], b[:, j, :]], axis=1))

有一个similar question使用“新单位尺寸”，但我不知道如何将tf.concat与“新单位尺寸”一起使用。

Answer 1

您可以将tf.tile和tf.expand_dims与tf.concat一起使用。一个例子：

import tensorflow as tf

a = tf.random_normal(shape=(2,4,5),dtype=tf.float32)
b = tf.random_normal(shape=(2,3,5),dtype=tf.float32)

# your code
concat_list = []
for i in range(a.get_shape()[1]):
    for j in range(b.get_shape()[1]):
        concat_list.append(tf.concat([a[:, i, :], b[:, j, :]], axis=1))

# Application  method
A = tf.tile(tf.expand_dims(a,axis=1),[1,b.shape[1],1,1])
B = tf.tile(tf.expand_dims(b,axis=2),[1,1,a.shape[1],1])
result = tf.concat([A,B],axis=-1)

with tf.Session() as sess:
    concat_list_val,result_val = sess.run([concat_list,result])
    print(concat_list_val[-1])
    print(result_val.shape)
    print(result_val[:,-1,-1,:])

# your result
[[ 1.0459949   1.5562199  -0.04387079  0.17898582 -1.9795663   0.988437
  -0.40415847  0.8865694  -1.4764767  -0.8417388 ]
 [-0.3542176  -0.3281141   0.01491702  0.91899025 -1.0651684   0.12315683
   0.6555444  -0.80451876 -1.3260773   0.33680603]]
# Application result shape
(2, 3, 4, 10)
# Application result 
[[ 1.0459949   1.5562199  -0.04387079  0.17898582 -1.9795663   0.988437
  -0.40415847  0.8865694  -1.4764767  -0.8417388 ]
 [-0.3542176  -0.3281141   0.01491702  0.91899025 -1.0651684   0.12315683
   0.6555444  -0.80451876 -1.3260773   0.33680603]]

性能

您可以使用以下代码来比较速度。

import datetime
...

with tf.Session() as sess:
    start = datetime.datetime.now()
    print('#' * 60)
    for i in range(10000):
        result_val = sess.run(result)
    end = datetime.datetime.now()
    print('cost time(seconds) : %.2f' % ((end - start).total_seconds()))

    start = datetime.datetime.now()
    print('#' * 60)
    for i in range(10000):
        concat_list_val = sess.run(concat_list)
    end = datetime.datetime.now()
    print('cost time(seconds) : %.2f' % ((end - start).total_seconds()))

当我的8GB GPU内存上有1.48s和5.76s时，向量化方法10000次迭代需要a.shape=(2,4,5)，而循环10000次迭代需要b.shape=(2,3,5)。但是向量化方法需要3.28s，并且当317.23s和a.shape=(20,40,5)时，循环时间为b.shape=(20,40,5)。

向量化方法将比tf.map_fn()和python循环快得多。

张量流嵌套map_fn concat两个张量

1 个答案: