TensorFlow：如何切片/收集所有可能的配置？

Question

我有一个带有形状的张量（批处理大小，序列长度，2，n，k）（在我的特殊情况下，2表示（x，y）空间位置）。 n表示n变量，k是每个变量所占的值。

我想生成带有shape 的张量n“变量”采取了其每个k可能的值。

如何使用切片或聚集在张量流中有效地进行此操作？

让我提供一个示例。出于插图的目的，我将省略批处理大小和序列长度的两个领先维度。

假设这是形状的3D张量x（2，n = 3，k = 4）：

第一个配置将是（滥用稍微表示符号），将切片如x [：，（0，1，2），（0，0，0）];在这里，n = 3变量都采用其第一个值。第二个配置是将切片如x [：，（0，1，2），（0，0，1）]];在这里，前两个变量正在采用其第一个值，第三个变量正在采用其第二个值。这会持续到4^3 = 64可能的配置，最后一个为x [：，（0，1，2），（3，3，3）]。

如果我将所有这些都堆叠在一起，则结果将是带状（2，3，4^3）的张量。

Answer 1

iiuc，如果仍然相关，这是您可以用tensorflow纯粹解决问题的一种方法（请注意，我与3D张量一起使用并省略了前两个领先维度）：

import tensorflow as tf

tf.random.set_seed(111)
x = tf.random.uniform((2, 3, 4), maxval=15, dtype=tf.int32)
x_shape = tf.shape(x)
print('x -->', x, '\n')

first_dim = tf.range(x_shape[0])
second_dim = tf.range(x_shape[1])
second_dim = tf.repeat(second_dim, tf.shape(first_dim)[0])
combination_range = tf.range(x_shape[-1])
xx, yy, zz = tf.meshgrid(combination_range, combination_range, combination_range, indexing='ij')
combinations = tf.stack([tf.reshape(xx, [-1]), tf.reshape(yy, [-1]), tf.reshape(zz, [-1])], axis=1)
print('combinations -->', combinations, '\n')

combinations = tf.reshape(tf.tile(combinations, [1, tf.shape(first_dim)[0]]), [-1])
first_dim = tf.tile(first_dim, [tf.shape(combinations)[0] // tf.shape(first_dim)[0]])
second_dim = tf.tile(second_dim, [tf.shape(combinations)[0] // tf.shape(second_dim)[0]])

result = tf.gather_nd(x, tf.transpose(tf.stack([first_dim, second_dim, combinations])))
result = tf.reshape(result, (x_shape[0], x_shape[1], x_shape[-1]**x_shape[1]))
print('final result -->', result)

x --> tf.Tensor(
[[[ 5 14  1 14]
  [ 2  1 12  3]
  [ 2  5  7 10]]

 [[ 0  9  0 12]
  [12 11  0  1]
  [ 2  6  1 12]]], shape=(2, 3, 4), dtype=int32) 

combinations --> tf.Tensor(
[[0 0 0]
 [0 0 1]
 [0 0 2]
 [0 0 3]
 [0 1 0]
 [0 1 1]
 [0 1 2]
 [0 1 3]
 [0 2 0]
 [0 2 1]
 [0 2 2]
 [0 2 3]
 [0 3 0]
 [0 3 1]
 [0 3 2]
 [0 3 3]
 [1 0 0]
 [1 0 1]
 [1 0 2]
 [1 0 3]
 [1 1 0]
 [1 1 1]
 [1 1 2]
 [1 1 3]
 [1 2 0]
 [1 2 1]
 [1 2 2]
 [1 2 3]
 [1 3 0]
 [1 3 1]
 [1 3 2]
 [1 3 3]
 [2 0 0]
 [2 0 1]
 [2 0 2]
 [2 0 3]
 [2 1 0]
 [2 1 1]
 [2 1 2]
 [2 1 3]
 [2 2 0]
 [2 2 1]
 [2 2 2]
 [2 2 3]
 [2 3 0]
 [2 3 1]
 [2 3 2]
 [2 3 3]
 [3 0 0]
 [3 0 1]
 [3 0 2]
 [3 0 3]
 [3 1 0]
 [3 1 1]
 [3 1 2]
 [3 1 3]
 [3 2 0]
 [3 2 1]
 [3 2 2]
 [3 2 3]
 [3 3 0]
 [3 3 1]
 [3 3 2]
 [3 3 3]], shape=(64, 3), dtype=int32) 

final result --> tf.Tensor(
[[[ 5  0  2 12  2  2  5  0  1 12  2  6  5  0 12 12  2  1  5  0  3 12  2
   12  5  9  2 12  5  2  5  9  1 12  5  6  5  9 12 12  5  1  5  9  3 12
    5 12  5  0  2 12  7  2  5  0  1 12  7  6  5  0 12 12]
  [ 7  1  5  0  3 12  7 12  5 12  2 12 10  2  5 12  1 12 10  6  5 12 12
   12 10  1  5 12  3 12 10 12 14  0  2 11  2  2 14  0  1 11  2  6 14  0
   12 11  2  1 14  0  3 11  2 12 14  9  2 11  5  2 14  9]
  [ 1 11  5  6 14  9 12 11  5  1 14  9  3 11  5 12 14  0  2 11  7  2 14
    0  1 11  7  6 14  0 12 11  7  1 14  0  3 11  7 12 14 12  2 11 10  2
   14 12  1 11 10  6 14 12 12 11 10  1 14 12  3 11 10 12]]

 [[ 1  0  2  0  2  2  1  0  1  0  2  6  1  0 12  0  2  1  1  0  3  0  2
   12  1  9  2  0  5  2  1  9  1  0  5  6  1  9 12  0  5  1  1  9  3  0
    5 12  1  0  2  0  7  2  1  0  1  0  7  6  1  0 12  0]
  [ 7  1  1  0  3  0  7 12  1 12  2  0 10  2  1 12  1  0 10  6  1 12 12
    0 10  1  1 12  3  0 10 12 14  0  2  1  2  2 14  0  1  1  2  6 14  0
   12  1  2  1 14  0  3  1  2 12 14  9  2  1  5  2 14  9]
  [ 1  1  5  6 14  9 12  1  5  1 14  9  3  1  5 12 14  0  2  1  7  2 14
    0  1  1  7  6 14  0 12  1  7  1 14  0  3  1  7 12 14 12  2  1 10  2
   14 12  1  1 10  6 14 12 12  1 10  1 14 12  3  1 10 12]]], shape=(2, 3, 64), dtype=int32)

对于arbitary n <n < /代码>，尝试以下操作：

import tensorflow as tf

tf.random.set_seed(111)
x = tf.random.uniform((2, 6, 4), maxval=15, dtype=tf.int32)
x_shape = tf.shape(x)
print('x -->', x, '\n')

first_dim = tf.range(x_shape[0])
second_dim = tf.range(x_shape[1])
second_dim = tf.repeat(second_dim, tf.shape(first_dim)[0])
combination_range = tf.range(x_shape[-1])
outputs = tf.meshgrid(*tuple(tf.unstack(tf.repeat(combination_range[tf.newaxis, ...], x_shape[1], axis=0), axis=0)), indexing='ij')
combinations = tf.stack([tf.reshape(o, [-1]) for o in outputs], axis=1)
print('combinations -->', combinations, '\n')

combinations = tf.reshape(tf.tile(combinations, [1, tf.shape(first_dim)[0]]), [-1])
first_dim = tf.tile(first_dim, [tf.shape(combinations)[0] // tf.shape(first_dim)[0]])
second_dim = tf.tile(second_dim, [tf.shape(combinations)[0] // tf.shape(second_dim)[0]])

result = tf.gather_nd(x, tf.transpose(tf.stack([first_dim, second_dim, combinations])))
result = tf.reshape(result, (x_shape[0], x_shape[1], x_shape[-1]**x_shape[1]))
print('final result -->', result)