Как сделать тензор потока сегмент_макс в высоком измерении

Question

Я хочу иметь возможность вызывать tf.math.unsorted_segment_max tenorsflow на тензор данных, который имеет размер [N, s, K]. N - количество каналов, а K - количество фильтров / карт характеристик. s - размер выборки данных одного канала. У меня есть factor_ids в размере s. Например, скажем, мой размер выборки s = 6, и я хочу сделать максимум для двух элементов (как если бы я делал обычное максимальное объединение, так что во втором, s-измерение всего тензора данных). Тогда мой мой_сегмент равен [0,0,1,1,2,2].

Я пытался запустить

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

с расширенными 0 и 2 измерениями для plot_ids, но так как идентификаторы сегмента затем повторяются, результат, конечно, будет размером [3] вместо [N, 3, K], как мне бы хотелось.

Так что мой вопрос в том, как построить правильный тензор сегмент_идейсов, чтобы добиться того, чего я хочу? То есть сделать max сегмента на основе исходного тензора сегмента размера s, но в каждом измерении отдельно?

По сути, возвращаясь к примеру, учитывая список идентификаторов 1D сегмента seg_id = [0,0,1,1,2,2], я хотел бы построить что-то вроде тензора сегмента_данных, для которого:

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

Так что при вызове tf.math. (Unsorted_) сегмент_max с этим тензором в качестве идентификаторов сегмента я получу результат размера [N, 3, K], с тем же эффектом, как если бы для Каждые данные [x,:, y] отдельно и соответствующим образом складывать результаты.

Любой способ сделать это хорошо, если он работает с тензорным потоком. Я бы предположил, что комбинация tf.tile, tf.reshape или tf.concat должна сработать, но я не могу понять, как, в каком порядке. Кроме того, есть ли более простой способ сделать это? Без необходимости настраивать сегмент_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, вы можете сделать это, как вы предлагаете в свой ответ .Вот эквивалентная формулировка, которая избегает транспонирования и включает окончательное изменение формы:

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])

Сэто можно вызывать функцией plot_max напрямую:

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

И он тоже делает то, что я хочу, за исключением того, что результат сглаживается и требует повторного изменения формы, если это необходимо.Эквивалентно, вы можете преобразовать исходный тензор данных в 1d, и на нем вы можете вызвать Calle_max, указав a + b в качестве_сегмента.И снова измените окончательный результат, если это необходимо.

Это то, что кажется долгим путем к результату ... Есть ли лучший способ?Я также не знаю, подходит ли описанный способ для использования внутри NN, во время backprop ... могут ли быть проблемы с производными или графом вычислений?У кого-нибудь есть идеи о том, как решить эту проблему?