python之通过 Tensorflow 中的段长度计算 tf.math.segment_sum 中所需的段 ID

我正在处理可变大小的顺序数据。 让我们考虑像这样的数据

Y = [ [.01,.02], [.03,.04], [.05,.06], [.07,.08], [.09,.1] ] 
l = [ 3, 2 ] 

其中 Y 是对我的数据执行的一些辅助计算的结果，l 存储原始序列的长度。在此示例中，[.01,.02]、[.03,.04]、[.05,.06] 因此是对批处理的第一个序列执行计算的结果，并且 [.07,.08], [.09,.1] 是对长度为 3 和 2 的批处理的第二个序列执行的计算结果 分别。 现在我想对 Y 的条目做一些进一步的计算，但按序列分组。在 Tensorflow 中，有 tf.math.segment_sum 等函数，可以按组执行。

假设我想使用 tf.math.segment_sum 求和。我会感兴趣

seq_ids = [ 0, 0, 0, 1, 1 ] 
tf.math.segment_sum(Y, segment_ids=seq_ids) #returns [ [0.09 0.12], [0.16 0.18] ] 

我现在面临的问题是从l 获取seq_ids。 在 numpy 中，可以很容易地通过

seq_ids = np.digitize( np.arange(np.sum(l)), np.cumsum(l) ) 

似乎有一个隐藏的(来自 python api)等效于 digitize 名为 bucketize，如 this search 中所述对于 Tensorflow 中的 digitize。 但似乎引用的 hidden_​​ops.txt 已从 Tensorflow 中删除，我不清楚是否仍然(并且将会)支持函数 tensorflow::ops::Bucketize 在 python api 中。 我必须获得类似结果的另一个想法是使用 tf.train.piecewise_constant 函数。但是这次尝试失败了，因为

seq_ids = tf.train.piecewise_constant(tf.range(tf.math.reduce_sum(l)), tf.math.cumsum(l), tf.range(BATCH_SIZE-1)) 

“Tensor”类型的对象没有 len() 失败。 似乎 tf.train.piecewise_constant 没有以最通用的方式实现，因为参数 boundaries 和 values 需要是列表而不是张量。作为 l 在我的例子中是一个一维张量，收集在我的 tf.data.Dataset

请您参考如下方法:

这是一种方法:

import tensorflow as tf 
 
def make_seq_ids(lens): 
    # Get accumulated sums (e.g. [2, 3, 1] -> [2, 5, 6]) 
    c = tf.cumsum(lens) 
    # Take all but the last accumulated sum value as indices 
    idx = c[:-1] 
    # Put ones on every index 
    s = tf.scatter_nd(tf.expand_dims(idx, 1), tf.ones_like(idx), [c[-1]]) 
    # Use accumulated sums to generate ids for every segment 
    return tf.cumsum(s) 
 
with tf.Graph().as_default(), tf.Session() as sess: 
    print(sess.run(make_seq_ids([2, 3, 1]))) 
    # [0 0 1 1 1 2] 

编辑:

您也可以使用 tf.searchsorted 实现相同的功能，与您为 NumPy 提出的建议更相似:

import tensorflow as tf 
 
def make_seq_ids(lens): 
    c = tf.cumsum(lens) 
    return tf.searchsorted(c, tf.range(c[-1]), side='right') 

这两种实现都不应成为 TensorFlow 模型的瓶颈，因此对于任何实际目的而言，您选择哪一种都无关紧要。然而，有趣的是，在我的特定机器(Win 10、TF 1.12、Core i7 7700K、Titan V)中，第二个实现在 CPU 上运行时慢 1.5 倍，在 GPU 上运行时快 3.5 倍。