给定 TF2.1 中每一行的起始索引，如何对张量进行切片？答案

【问题标题】：How to slice a Tensor given the starting indices for each each row in TF2.1?给定 TF2.1 中每一行的起始索引，如何对张量进行切片？
【发布时间】：2020-04-13 12:33:00
【问题描述】：

给定一些（至少是二维的）输入，例如：

inputs = [['a0', 'a1', 'a2', 'a3', 'a4'],
          ['b0', 'b1', 'b2', 'b3', 'b4'],
          ['c0', 'c1', 'c2', 'c3', 'c4']]

...以及另一个索引输入和标量窗口大小：

indices = [2, 3, 0]  # representing the starting positions (2nd dimension)
window_size = 2      # fixed-width of each window

如何在 Tensorflow 2 中获取从这些索引开始的窗口？我首先考虑使用像inputs[,start:start+window_size] 这样的常规切片，但这并不适用，因为这仅允许对所有行使用一个起始索引并且不支持每行不同的索引。

此示例的预期输出为：

output = [['a2', 'a3'],
          ['b3', 'b4'],
          ['c0', 'c1']]

【问题讨论】：

也许删除指定索引处的项目并将其保存到新列表中？ saveditem = inputs[0][indicies[0]] 然后inputs[0].remove(indices[0]) 然后newlist[0].append(saveditem)?
目前，我正在使用tf.map_fn 遍历所有行并将它们单独切片。这可行，但似乎很奇怪，因为我什至不确定这是否会在 GPU 上并行执行，或者在自定义网络层中使用时是否会导致瓶颈。
slicing 比 for 循环更有效，并且更具可读性，这是 python，与 C++ 相比，python 无论如何都慢。一个好的序列化程序应该采用它给定的任何结构并准备好发送，只要在发送到网络之前所有内容都经过预处理就可以了，对吧？

标签： python tensorflow tensorflow2.0

【解决方案1】：

我提供了一种矢量化方法。向量化方法将明显快于tf.map_fn()。

import tensorflow as tf

inputs = tf.constant([['a0', 'a1', 'a2', 'a3', 'a4'],
                      ['b0', 'b1', 'b2', 'b3', 'b4'],
                      ['c0', 'c1', 'c2', 'c3', 'c4']])
indicies = tf.constant([2, 3, 0])
window_size = 2

start_index = tf.sequence_mask(indicies,inputs.shape[-1])
# tf.Tensor(
# [[ True  True False False False]
#  [ True  True  True False False]
#  [False False False False False]], shape=(3, 5), dtype=bool)
end_index = tf.sequence_mask(indicies+window_size,inputs.shape[-1])
# tf.Tensor(
# [[ True  True  True  True False]
#  [ True  True  True  True  True]
#  [ True  True False False False]], shape=(3, 5), dtype=bool)

index = tf.not_equal(start_index,end_index)
# tf.Tensor(
# [[False False  True  True False]
#  [False False False  True  True]
#  [ True  True False False False]], shape=(3, 5), dtype=bool)

result = tf.reshape(tf.boolean_mask(inputs,index),
                    indicies.get_shape().as_list()+[window_size])
print(result)
# tf.Tensor(
# [[b'a2' b'a3']
#  [b'b3' b'b4']
#  [b'c0' b'c1']], shape=(3, 2), dtype=string)

【讨论】：