замена "tf.gather_nd" - PullRequest
       21

замена "tf.gather_nd"

1 голос
/ 05 июня 2019

Я делаю проект, но его версия tenorflow не поддерживает tf.gather_nd. Я спрашиваю, если возможно, что использовать tf.gather, tf.slice или tf.strided_slice, чтобы переписать функцию tf.gather_nd?

tf.gather_nd используется для сбора фрагментов из тензора в тензор с формой, определяемой индексами. подробности можно найти в https://www.tensorflow.org/api_docs/python/tf/gather_nd

Спасибо

1 Ответ

0 голосов
/ 05 июня 2019

Эта функция должна выполнять эквивалентную работу:

import tensorflow as tf
import numpy as np

def my_gather_nd(params, indices):
    idx_shape = tf.shape(indices)
    params_shape = tf.shape(params)
    idx_dims = idx_shape[-1]
    gather_shape = params_shape[idx_dims:]
    params_flat = tf.reshape(params, tf.concat([[-1], gather_shape], axis=0))
    axis_step = tf.cumprod(params_shape[:idx_dims], exclusive=True, reverse=True)
    indices_flat = tf.reduce_sum(indices * axis_step, axis=-1)
    result_flat = tf.gather(params_flat, indices_flat)
    return tf.reshape(result_flat, tf.concat([idx_shape[:-1], gather_shape], axis=0))

# Test
np.random.seed(0)
with tf.Graph().as_default(), tf.Session() as sess:
    params = tf.constant(np.random.rand(10, 20, 30).astype(np.float32))
    indices = tf.constant(np.stack([np.random.randint(10, size=(5, 8)),
                                    np.random.randint(20, size=(5, 8))], axis=-1))
    result1, result2 = sess.run((tf.gather_nd(params, indices),
                                 my_gather_nd(params, indices)))
    print(np.allclose(result1, result2))
    # True
Добро пожаловать на сайт PullRequest, где вы можете задавать вопросы и получать ответы от других членов сообщества.
...