python - 如何在高维中执行tensorflow segment_max

Question

我希望能够调用tensorflow的tf.math.unsorted_segment_max在大小为 [N, s, K] 的数据张量上。 N 是 channel 数，K 是滤波器/特征图的数量。s 是单 channel 数据样本的大小。我的segment_ids 大小为s。例如，假设我的样本大小为 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张量，以实现我想要的？ IE。根据原始 s 大小的segment_ids 张量来完成最大段，但在每个维度上分别进行？

基本上，回到这个例子，给定 1D 段 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 相同每个数据[x,:,y] 分开并适本地堆叠结果。

任何方法都可以，只要它适用于 tensorflow 。我猜想 tf.tile、tf.reshape 或 tf.concat 的组合应该可以解决问题，但我不知道如何、按什么顺序。另外，有没有更直接的方法来做到这一点？不需要在每个“池化”步骤中调整segment_ids？

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 中的建议进行操作。这是避免转置并包括最终 reshape 的等效公式:

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)