如何使用在另一个张量中指定的索引在 tensorflow 中创建一个中间范围仅设置为 True 的布尔掩码

Question

我有一个形状为 [None, 1] 的张量，由每个批次的值组成。使用这个张量，我必须创建一个布尔掩码，其值从相应批次的索引值（从张量）开始设置为 true 到固定长度，其余全部设置为 false。

例如，考虑索引张量是

[[1],[2],[2]]

假设每批所需的时间步长为 5，固定长度为 2，那么对于第一批，在从 1 开始到 2 结束的索引中（固定长度 = 2），值将是设置为真。其他批次也是如此。即，我希望创建的布尔掩码是

[[False,True,True,False,False],
 [False,False,True,True,False],
 [False,False,True,True,False]]

如何在不必为每批单独做的情况下实现上述目标？最好不要在张量流中使用参差不齐的特征？

index < tf.range(number_of_timesteps) 

上面可以用于极端设置True，但我找不到中间设置True的方法。

Answer 1

可以结合使用 tf.one_hot 和 tf.roll 来解决您的示例。

input = [[1],[2],[2]]
intermediate_output = tf.one_hot(input, 5)
output = intermediate_output + tf.roll(intermediate_output, shift=1,axis=2)

如果您需要从零和一转换为布尔值

output = tf.where(tf.equal(output, 1), True, False)

说明： tf.one_hot 用于将您的索引转换为中间的 one hot 表示。接下来 tf.roll 将中间表示移动 1。将其与中间表示结合并转换为布尔值会返回您想要的输出。

编辑：

除了 for 循环之外，我没有看到将其扩展到多个时间步的好方法。下面的代码将生成您想要的输出

input = [[1],[2],[2]]
intermediate_output = tf.one_hot(input, 5)

outputs = []
timesteps = 3
for i in range(timesteps):
  outputs.append(tf.roll(intermediate_output, shift=i,axis=2))

output = tf.reduce_sum(tf.unstack(outputs, timesteps), 2)
output = tf.where(tf.equal(output, 1), True, False)