[英]Better way to access individual elements in a tensor
我正在嘗試使用張量b
定義的索引訪問張量a
的元素。
a=tf.constant([[1,2,3,4],[5,6,7,8]])
b=tf.constant([0,1,1,0])
我希望輸出是
out = [1 6 7 4]
我嘗試過的
out=[]
for i in range(a.shape[1]):
out.append(a[b[i],i])
out=tf.stack(out) #[1 6 7 4]
這給出了正確的輸出,但是我正在尋找一種更好,更緊湊的方式來實現。
同時我的邏輯不工作時的造型a
是一樣的東西(2,None)
,因為我不能重復range(a.shape[1])
它將如果答案包含在此情況下,也幫助我
謝謝
您可以使用tf.one_hot()
和tf.boolean_mask()
。
import tensorflow as tf
import numpy as np
a_tf = tf.placeholder(shape=(2,None),dtype=tf.int32)
b_tf = tf.placeholder(shape=(None,),dtype=tf.int32)
index = tf.one_hot(b_tf,a_tf.shape[0])
out = tf.boolean_mask(tf.transpose(a_tf),index)
a=np.array([[1,2,3,4],[5,6,7,8]])
b=np.array([0,1,1,0])
with tf.Session() as sess:
print(sess.run(out,feed_dict={a_tf:a,b_tf:b}))
# print
[1 6 7 4]
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.