使用另一個張量進行Tensorflow多維索引

Question

說，我有一個形狀為[n1，n2，...，nk]的秩k張量X和一個形狀為[n2，n3，...，nk]的秩（k-1）張量IDX，其中IDX具有與X的最后（k-1）維相同的形狀。IDX的條目都是[0，n1）中的所有整數。 我想從X提取一些值，其中第一個維度位置由IDX指定，而其他維度則一直迭代。

例：

X = tf.constant([[1,2], [3,4], [5,6],
                 [7,8], [9,10],[11,12]]) # 2 x 3 x 2 tensor
IDX = tf.constant([[1,0], [1,1], [0,1]]) #     3 x 2 tensor
...
# would like to get [[7,2],[9,10],[5,12]]

如何在Tensorflow中有效實現這一目標？ 謝謝！

Answer 1

您是否看到choose的注釋？

筆記

為了減少錯誤解釋的機會，即使名義上支持以下“濫用”， choices也不應該也不應被視為單個數組，即，最外面的類似於序列的容器應該是列表或元組。

也就是說，他們希望您將其視為：

In [432]: list(X)
Out[432]: [array([1, 2]), array([3, 4]), array([5, 6])]
In [433]: np.choose(IDX,list(X))
Out[433]: array([3, 6])

等效的索引是：

In [436]: X[IDX,np.arange(2)]
Out[436]: array([3, 6])

choose還具有一些mode選項。

文檔還說這等效於（減去這些模式問題）：

np.choose(a,c) == np.array([c[a[I]][I] for I in ndi.ndindex(a.shape)])

choose另一個細微差別。 它不能使用超過32個選擇。

In [440]: np.choose(IDX,np.ones((33,2)))
...
ValueError: Need at least 1 and at most 32 array objects.

In [442]: np.ones((33,2))[IDX,np.arange(2)]
Out[442]: array([ 1.,  1.])

Answer 2

您可以將np.choose()包裝在python函數中，並使用tf.py_func()將其嵌入到張量流圖中。 但是，如果您希望對圖進行自動梯度計算以進行訓練，則還可以為函數定義梯度。 我想，為np.choose()定義梯度可能是非常棘手的任務，即使實際上可以解決的話。