如何通过沿某个轴交错来合并两个 3D 张量？答案

【问题标题】：How can I merge two 3D tensors by interleaving them along a certain axis?如何通过沿某个轴交错来合并两个 3D 张量？
【发布时间】：2019-07-04 12:41:42
【问题描述】：

假设我想通过沿一个特定轴交错来合并两个 3D 张量流张量 a 和 b。例如，张量 a 的形状为 (3,3,2)，张量 b 的形状为 (3,2,2)。我想创建一个沿轴 1 交错的张量 c，从而产生一个形状为 (3,5,2) 的张量。

例子：

a = [[[1,1],[2,2],[3,3]],
     [[4,4],[5,5],[6,6]],
     [[7,7],[8,8],[9,9]]]

b = [[[10,10],[11,11]],
     [[12,12],[13,13]],
     [[14,14],[15,15]]]

c = [[[1,1],[10,10],[2,2],[11,11],[3,3]],
     [[4,4],[12,12],[5,5],[13,13],[6,6]],
     [[7,7],[14,14],[8,8],[15,15],[9,9]]]

【问题讨论】：

标签： python tensorflow

【解决方案1】：

您可以先对列的索引重新排序。

import tensorflow as tf

a = [[[1,1],[2,2],[3,3]],
     [[4,4],[5,5],[6,6]],
     [[7,7],[8,8],[9,9]]]

b = [[[10,10],[11,11]],
     [[12,12],[13,13]],
     [[14,14],[15,15]]]

a_tf = tf.constant(a)
b_tf = tf.constant(b)

a_tf_column = tf.range(a_tf.shape[1])*2
# [0 2 4]
b_tf_column = tf.range(b_tf.shape[1])*2+1
# [1 3]

column_indices = tf.concat([a_tf_column,b_tf_column],axis=-1)
# Before TF v1.13
column_indices = tf.contrib.framework.argsort(column_indices)
## From TF v1.13
# column_indices = tf.argsort(column_indices)

# [0 3 1 4 2]

那么您应该为tf.gather_nd() 创建新索引。

column,row = tf.meshgrid(column_indices,tf.range(a_tf.shape[0]))
combine_indices = tf.stack([row,column],axis=-1)
# [[[0,0],[0,3],[0,1],[0,4],[0,2]],
#  [[1,0],[1,3],[1,1],[1,4],[1,2]],
#  [[2,0],[2,3],[2,1],[2,4],[2,2]]]

最后你应该连接a和b的值并使用tf.gather_nd()得到结果。

combine_value = tf.concat([a_tf,b_tf],axis=1)
result = tf.gather_nd(combine_value,combine_indices)

with tf.Session() as sess:
    print(sess.run(result))

# [[[1,1],[10,10],[2,2],[11,11],[3,3]],
#  [[4,4],[12,12],[5,5],[13,13],[6,6]],
#  [[7,7],[14,14],[8,8],[15,15],[9,9]]]

【讨论】：

【解决方案2】：

这是我发现的一种更简单的交错方法：

import tensorflow as tf
a = tf.constant([1,2,3])
b = tf.constant([9,8,7])

# to interleave on axis N, stack on axis N+1
c = tf.stack([a,b], axis=1)
# reshape to remove the temp N+1 axis
c = tf.reshape(c, (-1,))

print(c)
# Output: [1, 9, 2, 8, 3, 7]

对于 OP 的示例，tf.stack 不能开箱即用，因为 A 有 3 个元素要交错，但 B 只有 2 个元素。我们可以通过将 B 填充到 3 个元素来轻松解决这个问题：

import tensorflow as tf
a = tf.constant([
    [[1,1],[2,2],[3,3]],
    [[4,4],[5,5],[6,6]],
    [[7,7],[8,8],[9,9]]
])
b = tf.constant([
    [[10,10],[11,11]],
    [[12,12],[13,13]],
    [[14,14],[15,15]]
])

# pad b to match a
b = tf.pad(b, [[0,0],[0,1],[0,0]])
# interleave on axis 1
c = tf.stack([a,b], axis=2)
c = tf.reshape(c, (a.shape[0],-1,a.shape[-1]))
# discard the padding elements
c = c[:,:-1]

【讨论】：