使用 tf.gather 逐行提取基于另一个张量的张量（第一维）答案

【问题标题】：Use tf.gather to extract tensors row-wise based on another tensor row-wisely (first dimension)使用 tf.gather 逐行提取基于另一个张量的张量（第一维）
【发布时间】：2019-07-23 17:07:10
【问题描述】：

我有两个尺寸分别为A:[B,3000,3] 和C:[B,4000] 的张量。我想使用 tf.gather() 将张量 C 中的每一行用作索引，并将张量 A 中的每一行用作参数，以获得大小为 [B,4000,3] 的结果。

这是一个更容易理解的例子：假设我有张量

A = [[1,2,3],[4,5,6],[7,8,9]],

C = [0,2,1,2,1],

结果 = [[1,2,3],[7,8,9],[4,5,6],[7,8,9],[4,5,6]],

使用tf.gather(A,C)。应用于维度小于 3 的张量时一切正常。

但是当以描述为开头的情况下，通过应用tf.gather(A,C,axis=1)，结果张量的形状是

[B,B,4000,3]

似乎tf.gather() 只是将张量 C 中的每个元素作为索引来收集张量 A 中的元素。我正在考虑的唯一解决方案是使用 for 循环，但这会极大地减少使用tf.gather(A[i,...],C[i,...])获得正确大小的张量的计算能力

[B,4000,3]

那么，是否有任何函数能够类似地完成这项任务？

【问题讨论】：

标签： python tensorflow tensor

【解决方案1】：

你需要使用tf.gather_nd:

import tensorflow as tf

A = ...  # B x 3000 x 3
C = ...  # B x 4000
s = tf.shape(C)
B, cols = s[0], s[1]
# Make indices for first dimension
idx = tf.tile(tf.expand_dims(tf.range(B, dtype=C.dtype), 1), [1, cols])
# Complete index for gather_nd
gather_idx = tf.stack([idx, C], axis=-1)
# Gather result
result = tf.gather_nd(A, gather_idx)

【讨论】：

这是一个很好的解决方案，谢谢！这对我来说有点难以理解。如果我错了，请告诉我。 'idx' 为每行（第一维）提供索引，'gather_idx' 为每个 'idx' 中的 4000 个元素提供索引，这样 'Bx4000x2' 为 A? 提供整个索引'gather_idx' (Bx4000x2) 中的第三个维度的实际含义是什么？是指B中的元素ID和4000中的元素ID所代表的全局索引吗？
@KayeeJoe 您可以将A 视为3 个向量的集合，每个向量都位于张量中的某个特定坐标(i, j)。在tf.gather_nd 中，gather_idx 的最后一个维度表示来自A 的这两个坐标。因此，如果gather_idx 的形状为(B, 4000, 2)，则意味着您将获得一个形状为(B, 4000, 3) 的“3 向量矩阵”，这样结果中(r, s) 位置处的3 向量将为A[gather_idx[r, s, 0], gather_idx[r, s, 1]]。在您的情况下，输入 C 只为您提供第二个索引，而使用 idx 您可以获得tf.gather_nd 的完整二维索引。