Ручная реализация функций unsorted_segment в Tensorflow - PullRequest
Новая
       1

Ручная реализация функций unsorted_segment в Tensorflow

0 голосов
/ 21 июня 2020

У меня есть код, в котором я использую unsorted_segment_mean, но хотел бы сделать то же самое с помощью медианной функции Tensorflow Probability. Это не реализовано, и, посмотрев на исходный код для unsorted_segment_mean, я не в состоянии реализовать.

Мне было интересно, есть ли более общий c, ручной способ реализации unsorted_segment_mean, даже если это происходит за счет небольшой стоимости производительности, которую я мог бы затем изменить для использования любой функции сокращения (например, медианы).

Моя попытка заключалась в следующем: 

def unsorted_segment(func, data, segment_ids, k):
    # `data` is nxp, `segment_ids` is nx1 with values in 0,...,k-1
    # `res` should be kxp
    res = tf.zeros((k, data.get_shape()[1]))
    for i in range(k):
        in_segment = tf.math.equal(segment_ids, i)
        if tf.math.reduce_any(in_segment):
            res[i] = tf.math.reduce_mean(data[is_segment], axis=0)
        else:
            # Special case to avoid division by zero...ignore this
            pass

Проблема с этим решением заключается в том, что назначение срезов (res[i] = ...) не разрешено в Tensorflow, и я не мог найти обходной путь для этого. Вдобавок к этому в строке math.equal много избыточности, поскольку она проверяет каждый элемент k раз. Можно ли улучшить / исправить эти проблемы или есть другой подход?

...