TensorFlow getting elements of every row for specific columns

后端 未结 5 702
再見小時候
再見小時候 2021-02-08 02:57

If A is a TensorFlow variable like so

A = tf.Variable([[1, 2], [3, 4]])

and index is another variable



        
相关标签:
5条回答
  • 2021-02-08 03:40

    You can extend your column indices with row indices and then use gather_nd:

    import tensorflow as tf
    
    A = tf.constant([[1, 2], [3, 4]])
    indices = tf.constant([1, 0])
    
    # prepare row indices
    row_indices = tf.range(tf.shape(indices)[0])
    
    # zip row indices with column indices
    full_indices = tf.stack([row_indices, indices], axis=1)
    
    # retrieve values by indices
    S = tf.gather_nd(A, full_indices)
    
    session = tf.InteractiveSession()
    session.run(S)
    
    0 讨论(0)
  • 2021-02-08 03:44

    You can use one hot method to create a one_hot array and use it as a boolean mask to select the indices you'd like.

    A = tf.Variable([[1, 2], [3, 4]])
    index = tf.Variable([0, 1])
    
    one_hot_mask = tf.one_hot(index, A.shape[1], on_value = True, off_value = False, dtype = tf.bool)
    output = tf.boolean_mask(A, one_hot_mask)
    
    0 讨论(0)
  • 2021-02-08 03:45

    After dabbling around for quite a while. I found two functions that could be useful.

    One is tf.gather_nd() which might be useful if you can produce a tensor of the form [[0, 0], [1, 1]] and thereby you could do

    index = tf.constant([[0, 0], [1, 1]])

    tf.gather_nd(A, index)

    If you are unable to produce a vector of the form [[0, 0], [1, 1]](I couldn't produce this as the number of rows in my case was dependent on a placeholder) for some reason then the work around I found is to use the tf.py_func(). Here is an example code on how this can be done

    import tensorflow as tf 
    import numpy as np 
    
    def index_along_every_row(array, index):
        N, _ = array.shape 
        return array[np.arange(N), index]
    
    a = tf.Variable([[1, 2], [3, 4]], dtype=tf.int32)
    index = tf.Variable([0, 1], dtype=tf.int32)
    a_slice_op = tf.py_func(index_along_every_row, [a, index], [tf.int32])[0]
    session = tf.InteractiveSession()
    
    a.initializer.run()
    index.initializer.run()
    a_slice = a_slice_op.eval() 
    

    a_slice will be a numpy array [1, 4]

    0 讨论(0)
  • 2021-02-08 03:58

    We can do the same using this combination of map_fn and gather_nd.

    def get_element(a, indices):
        """
        Outputs (ith element of indices) from (ith row of a)
        """
        return tf.map_fn(lambda x: tf.gather_nd(x[0], x[1]), 
                                      (a, indices),
                                      dtype = tf.float32)
    

    Here's an example usage.

    A = tf.constant(np.array([[1,2,3],
                              [4,5,6],
                              [7,8,9]], dtype = np.float32))
    
    idx = tf.constant(np.array([[2],[1],[0]]))
    elems = get_element(A, idx)
    
    with tf.Session() as sess:
        e = sess.run(elems)
    
    print(e)
    

    I don't know if this will be much slower than other answers.

    It has the advantage that you don't need to specify the number of rows of A in advance, as long as a and indices have the same number of rows at runtime.

    Note the output of the above will be rank 1. If you'd prefer it to have rank 2, replace gather_nd by gather

    0 讨论(0)
  • I couldn't get the accepted answer to work in Tensorflow 2 when I incorporated it into a loss function. Something about GradientTape didn't like it. My solution is an altered version of the accepted answer:

    def get_rows(arr):
      N, _ = arr.shape 
      return N
    
    num_rows= tf.py_function(get_rows, [arr], [tf.int32])[0]
    rng = tf.range(0,num_rows)
    ind = tf.stack([rng, ind], axis=1)
    tf.gather_nd(arr, ind)
    
    0 讨论(0)
提交回复
热议问题