我需要根據(jù)索引張量從張量中提取值。我的代碼如下：arr = tf.constant([10, 11, 12]) # array of valuesinds = tf.constant([0, 1, 2]) # indicesres = tf.map_fn(fn=lambda t: arr[t], elems=inds)它運作緩慢。有更有效的方法嗎？
1 回答
慕蓋茨4494581
TA貢獻1850條經(jīng)驗 獲得超11個贊
您可以使用 tf.gather 方法
arr = tf.constant([10, 11, 12]) # array of values
inds = tf.constant([0, 2])
r = tf.gather(arr , inds)#<tf.Tensor: shape=(2,), dtype=int32, numpy=array([10, 12])>
如果您有一個多維張量，則 tf.gather 有一個“axis”參數(shù)來指定檢查索引的維度：
arr = tf.constant([[10, 11, 12] ,[1, 2, 3]]) # shape(2,3)
inds = tf.constant([0, 1])
# axis == 1
r = tf.gather(arr , inds , axis = 1)#<tf.Tensor: shape=(2, 2), dtype=int32, numpy=array([[10, 11],[ 1, 2]])>
# axis == 0
r = tf.gather(arr , inds , axis = 0) #<tf.Tensor: shape=(2, 3), dtype=int32, numpy=array([[10, 11, 12], [ 1, 2, 3]])>
添加回答
舉報
0/150
提交
取消