挑选出tensor中等于0的索引_tensorflow实现tensor序列不定长截取

应用场景：当我们在实现深度学习模型时，由于模型内部是矩阵运算，所以要求一个batch中的每个样本的序列长度是相同的（不一样长要进行padding补0操作）；但是有时，我们需要获取一个batch样本中，每一个样本的不同长度的序列，进行其他操作。

举例：假设有一个batch_size为3，长度为8的二维Tensor-tensor_a，如下：

[[1 3 5 7 8 1 4 6]
 [2 1 7 4 6 1 4 9]
 [3 5 4 2 3 2 2 5]]

知道，需要截取的每个样本的起始位置标号start和结束位置标号end，如下：

start [3 1 2]
end [6 3 7]

想要获取这个二维tensor的不定长序列：

第一个样本 截取第3到第6个元素[3 5 7 8]
第二个样本 截取第1到第3个元素[2 1 7]
第三个样本 截取第2到第7个元素[5 4 2 3 2 2]

由于矩阵中要求每行要有相同个数的元素，因此我们通常采用掩码来解决这种问题。掩码，即是一个与原始矩阵维度一致的[0,1]矩阵，0表示对应位置被舍弃，1表示对应位置被保留。然后通过掩码与原始矩阵相乘，就得到了我们想要截取出的不定长序列。

具体操作如下：

【1】构建上述tensor_a，start和end矩阵：

tensor_a 

tensor_a，start和end的输出为：

[[1 3 5 7 8 1 4 6]
 [2 1 7 4 6 1 4 9]
 [3 5 4 2 3 2 2 5]]
[3 1 2]
[6 3 7]

【2】对start和end矩阵进行扩维操作，生成维度与tensor_a维度一致的矩阵，length 为tensor_a的序列长度：

length 

start和end的输出为：

[[3 3 3 3 3 3 3 3]
 [1 1 1 1 1 1 1 1]
 [2 2 2 2 2 2 2 2]]
[[6 6 6 6 6 6 6 6]
 [3 3 3 3 3 3 3 3]
 [7 7 7 7 7 7 7 7]]

【3】生成与原始矩阵维度一致的，元素值在1到length+1范围的序列矩阵，用于位置筛选，其中，元素值表示所在位置标号，bacth_size为样本个数：

bacth_size 

range_mat 的输出为：

[[

【4】将序列矩阵range_mat 与对起始位置矩阵start和结束位置矩阵end分别进行比较，获取mask_start和mask_end，并对其取“位与”得到最终掩码：

mask_start 

mask_start 、mask_end 和mask输出为：

[[

【5】将掩码与原始矩阵相乘，得到我们需要的序列：

need_seq = tf.multiply(tensor_a, mask)

need_seq的输出为：

[[0 0 5 7 8 1 0 0]
 [2 1 7 0 0 0 0 0]
 [0 5 4 2 3 2 2 0]]

最终获取到了所需要的不定长序列，大功告成！如果大家有什么更加简洁高效地方法，请留言！

---

如果有对文本蕴含或者机器阅读理解感兴趣的同学，可以关注我的另一个专栏：

自然语言处理相关论文

刘聪NLP：文本蕴含之孪生网络（Siamese Network）

刘聪NLP：论文阅读笔记：文本蕴含之ESIM

刘聪NLP：论文阅读笔记：文本蕴含之DIIN

刘聪NLP：论文阅读笔记：文本蕴含之BiMPM

刘聪NLP：机器阅读理解之DuReader数据集描述