tensorflow nested map_fn concat two tensors - tensorflow

Question

say I have two tensors:
a=Tensor("zeros_3:0", shape=(2, 4, 5), dtype=float32)
b=Tensor("ones_3:0", shape=(2, 3, 5), dtype=float32)
how can I concat each element along axis 2 to get a new tensor shaped (2,3,4,10), using nested map_fn or other tf functions?
here is my for loop version
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))
there is a similar question using "new unit dimension", but I don't know how to use tf.concat with the "new unit dimension".

Answer 1

You can use tf.tile and tf.expand_dims with tf.concat. An example:
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]]
Performance
You can use follow code to compare speed.
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()))
The vectorization method 10000 iterations takes 1.48s and the loop 10000 iterations takes 5.76s when a.shape=(2,4,5) and b.shape=(2,3,5) on my 8GB GPU memory. But the vectorization method takes 3.28s and the loop time is 317.23s when a.shape=(20,40,5) and b.shape=(20,40,5).
The vectorization method will be significantly faster than the tf.map_fn() and python loop.