Question

我希望能够调用tensorflow的tf.math.unsorted_segment_max 在大小为[N，s，K]的数据张量上。 N是通道数，K是过滤器/功能图的数量。 s是单通道数据样本的大小。我有s个大小的segment_ids。例如，假设我的样本大小为s = 6，并且我想对两个元素做一个最大值（就好像进行通常的最大池化一样，第二个是整个数据张量的s维）。然后我的segment_ids等于[0,0,1,1,2,2]。

我尝试跑步

tf.math.unsorted_segment_max(data, segment_ids, num_segments)

具有segment_ids的扩展0和2维，但是由于段id然后重复，因此结果当然是[3]大小，而不是我想要的[N，3，K]。

所以我的问题是，如何构建适当的segment_ids张量，以实现我想要的？即根据原始的s大小segment_ids张量完成线段最大值，但是在每个维度上分别？

基本上，回到示例，给定一维段ID列表seg_id = [0,0,1,1,2,2]，我想为其构建一个类似于segment_ids张量的东西：

segment_ids[i,:,j] = seg_id + num_segments*(i*K + j)

因此，当使用该张量作为段ID调用tf.math。（unsorted_）segment_max时，我将得到大小为[N，3，K]的结果，就像在运行segment_max用于每个data [x，：，y]分别并适当地堆叠结果。

任何一种这样做的方法都是可以的，只要它与tensorflow一起工作即可。我猜想tf.tile，tf.reshape或tf.concat的组合应该可以解决问题，但是我不知道如何，以什么顺序。另外，还有更简单的方法吗？无需在每个“合并”步骤中调整segment_id？

Answer 1

我认为您可以通过tf.nn.pool实现自己想要的目标：

import tensorflow as tf

with tf.Graph().as_default(), tf.Session() as sess:
    data = tf.constant([
        [
            [ 1, 12, 13],
            [ 2, 11, 14],
            [ 3, 10, 15],
            [ 4,  9, 16],
            [ 5,  8, 17],
            [ 6,  7, 18],
        ],
        [
            [19, 30, 31],
            [20, 29, 32],
            [21, 28, 33],
            [22, 27, 34],
            [23, 26, 35],
            [24, 25, 36],
        ]], dtype=tf.int32)
    segments = tf.constant([0, 0, 1, 1, 2, 2], dtype=tf.int32)
    pool = tf.nn.pool(data, [2], 'MAX', 'VALID', strides=[2])
    print(sess.run(pool))

输出：

[[[ 2 12 14]
  [ 4 10 16]
  [ 6  8 18]]

 [[20 30 32]
  [22 28 34]
  [24 26 36]]]

如果您确实想要我们tf.unsorted_segment_max，则可以按照your own answer中的建议进行操作。这是一种避免转置的等效公式，其中包括最终的重塑：

import tensorflow as tf

with tf.Graph().as_default(), tf.Session() as sess:
    data = ...
    segments = ...
    shape = tf.shape(data)
    n, k = shape[0], shape[2]
    m = tf.reduce_max(segments) + 1
    grid = tf.meshgrid(tf.range(n) * m * k,
                       segments * k,
                       tf.range(k), indexing='ij')
    segment_nd = tf.add_n(grid)
    segmented = tf.unsorted_segment_max(data, segment_nd, n * m * k)
    result = tf.reshape(segmented, [n, m, k])
    print(sess.run(result))
    # Same output

从反向传播角度来看，这两种方法在神经网络中应该都能正常工作。

编辑：就性能而言，池化似乎比分段总和更具可扩展性（正如人们所期望的那样）：

import tensorflow as tf
import numpy as np

def method_pool(data, window):
    return tf.nn.pool(data, [window], 'MAX', 'VALID', strides=[window])

def method_segment(data, window):
    shape = tf.shape(data)
    n, s, k = shape[0], shape[1], shape[2]
    segments = tf.range(s) // window
    m = tf.reduce_max(segments) + 1
    grid = tf.meshgrid(tf.range(n) * m * k,
                       segments * k,
                       tf.range(k), indexing='ij')
    segment_nd = tf.add_n(grid)
    segmented = tf.unsorted_segment_max(data, segment_nd, n * m * k)
    return tf.reshape(segmented, [n, m, k])

np.random.seed(100)
rand_data = np.random.rand(300, 500, 100)
window = 10
with tf.Graph().as_default(), tf.Session() as sess:
    data = tf.constant(rand_data, dtype=tf.float32)
    res_pool = method_pool(data, n)
    res_segment = method_segment(data, n)
    print(np.allclose(*sess.run([res_pool, res_segment])))
    # True
    %timeit sess.run(res_pool)
    # 2.56 ms ± 80.8 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
    %timeit sess.run(res_segment)
    # 514 ms ± 6.29 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answer 2

我没有想出更多优雅的解决方案，但至少我想出了如何结合平铺，重塑和转置来实现的方法。我首先（使用上述三个操作，请参见下面的代码）构造一个与数据大小相同的张量，并在张量中重复（但移位）原始seg_id向量的条目：

m = tf.reduce_max(seg_id) + 1
a = tf.constant([i*m for i in range(N*K) for j in range(s)])
b = tf.tile(seg_id, N*K)
#now reshape it:
segment_ids = tf.transpose(tf.reshape(a+b, shape=[N,K,s]), perm=[0,2,1])

通过这种方式，可以直接调用segment_max函数：

result = tf.unsorted_segment_max(data=data, segment_ids=segment_ids, num_segments=m*N*K)

它也做了我想要的，除了结果被弄平并且需要时，需要再次重塑形状。同样，您可以将原始数据张量整形为1d，并使用a + b作为segment_ids在其上计算segment_max。并根据需要重新调整最终结果。

感觉到结果还有很长的路要走...还有更好的方法吗？我也不知道所描述的方法是否适合在反向传播期间在NN内部使用...导数或计算图是否可能存在问题？有人对如何解决这个问题有更好的主意吗？

如何在高维中做tensorflow segment_max

2 个答案: