2016-05-04 28 views
6

को देखते हुए में पंक्ति प्रति एक तत्व का चयन करने के लिए रास्ता ...सुरुचिपूर्ण Tensorflow

  • एक मैट्रिक्स आकार [m, n]
  • एक टेन्सर आकार के I[m]

मैं एक सूची प्राप्त करना चाहते हैं की A से J तत्वों के J[i] = A[i, I[i]]

है, IA में प्रत्येक पंक्ति से चयन करने के लिए तत्व की अनुक्रमणिका रखता है।

संदर्भ: मेरे पास पहले से ही argmax(A, 1) है और अब मैं max भी चाहता हूं। मुझे पता है कि मैं सिर्फ reduce_max का उपयोग कर सकता हूं। और थोड़ा-बहुत प्रयास करने के बाद मैं भी इस के साथ आया था:

J = tf.gather_nd(A, 
    tf.transpose(tf.pack([tf.to_int64(tf.range(A.get_shape()[0])), I]))) 

कहाँ to_int64 की जरूरत है क्योंकि सीमा केवल int32 पैदा करता है और argmax केवल int64 पैदा करता है।

दोनों में से कोई भी विशेष रूप से सुरुचिपूर्ण के रूप में मुझे हड़ताल नहीं करता है। एक रनटाइम ओवरहेड है (शायद कारक n के बारे में) और दूसरे के पास अज्ञात कारक संज्ञानात्मक उपरि है। क्या मुझसे कोई चूक हो रही है?

def get_elements(data, indices): 
    indeces = tf.range(0, tf.shape(indices)[0])*data.shape[1] + indices 
    return tf.gather(tf.reshape(data, [-1]), indeces) 

आपका समाधान वर्तमान में जो विभेदक नहीं है (क्योंकि tf.gather_nd के लिए ढ़ाल समय समर्थित नहीं हैं):

+0

आप 'tf.range',' tf.pack' और 'tf का उपयोग करके सूची 'ind = [[1, I [1]], [2, I [2]], ...] बना सकते हैं। .transpose', और फिर 'tf.gather_nd (जे, इंड)' –

+0

@YaroslavBulatov बिल्कुल वही नहीं है जो मैंने वर्णित किया है? –

+0

लेकिन मुझे लगता है कि इसका मतलब यह होगा कि यह सिर्फ रास्ता है? –

उत्तर

0

Link @ यारोस्लाव-Bulatov द्वारा प्रदान की this समाधान mentiones।

उम्मीद है कि data[:, indices] जल्द ही पेश किया जाएगा।

2

यह एक नहीं बल्कि देर से जवाब है, लेकिन

mask = tf.one_hot(I, depth=n, dtype=tf.bool, on_value=True, off_value=False) 
elements = tf.boolean_mask(A, mask) 

पूरा आप जो खोज रहे हैं कर सकता है?

संपादित करें: मुझे यह इंगित करना चाहिए कि A पहले से ही एक बहुत बड़ा टेंसर है, क्योंकि यह एक घना मैट्रिक्स बनाने के समाप्त होता है।