тензорное многомерное индексирование

Question

У меня есть

• 4-мерный тензор с плавающей точкой y и

• 3-мерный целочисленный тензор y_index, содержащий индекс 4-е измерение y для извлечения

То, что я хочу сделать, довольно просто с 3 for-l oop с numpy:

y = np.random.randint(100,size=(5,10,20,3))
y_index= np.random.randint(3,size=(5,10,20))
y_slice = np.zeros_like(y_index)
for i in range(y.shape[0]):
    for j in range(y.shape[1]):
        for k in range(y.shape[2]):
            y_slice[i,j,k] = y[i,j,k,y_index[i,j,k]]
y_slice

Как я могу сделать это эффективно в tenorflow? Я догадываюсь, что мне нужно использовать tf.gether_nd ...

Answer 1

Вы можете сделать следующее. По сути, сначала вы сглаживаете все измерения, кроме последнего из y, и создаете индекс для сглаживания y. Вы выполняете индексацию, а затем изменяете форму до правильной формы.

y = tf.constant(np.random.normal(size=(5,10,20,3)), dtype='float32')
y_index = tf.constant(np.random.randint(3, size=(5,10,20)), dtype='int32')
# Creating an index like [(0,y_index[0]), (1, y_index[1]), ...]
inds = tf.stack([tf.range(5*10*20),tf.reshape(y_index,[-1])],axis=1)

y_slice = tf.reshape(tf.gather_nd(tf.reshape(y,[-1,3]),inds),[5,10,20])