How to do the tf.gather_nd in pytorch?

Alex_Hex: <p>In Tensorflow we can do the below operation by tf.gather_nd, but how to do this in pytorch?<br>
Simple indexing into a matrix:</p>
<pre><code>indices = [[0, 0], [1, 1]]
params = [['a', 'b'], ['c', 'd']]
output = ['a', 'd']
</code></pre>
<p>Slice indexing into a matrix:</p>
<pre><code>indices = [[1], [0]]
params = [['a', 'b'], ['c', 'd']]
output = [['c', 'd'], ['a', 'b']]
</code></pre>
<p>Indexing into a 3-tensor:</p>
<pre><code>indices = [[1]]
params = [[['a0', 'b0'], ['c0', 'd0']],
          [['a1', 'b1'], ['c1', 'd1']]]
output = [[['a1', 'b1'], ['c1', 'd1']]]

indices = [[0, 1], [1, 0]]
params = [[['a0', 'b0'], ['c0', 'd0']],
          [['a1', 'b1'], ['c1', 'd1']]]
output = [['c0', 'd0'], ['a1', 'b1']]

indices = [[0, 0, 1], [1, 0, 1]]
params = [[['a0', 'b0'], ['c0', 'd0']],
          [['a1', 'b1'], ['c1', 'd1']]]
output = ['b0', 'b1']
</code></pre>
<p>Batched indexing into a matrix:</p>
<pre><code>indices = [[[0, 0]], [[0, 1]]]
params = [['a', 'b'], ['c', 'd']]
output = [['a'], ['b']]</code></pre>

thor: <p>Hi. Did you ever discover the best way to deal with this issue?</p>

richard: <p>Yes, there are equivalent operations in pytorch. Try something like the following:</p>
<p>Simple indexing into matrix:</p>
<pre><code class="lang-auto">x = torch.randn(2, 2)
indices = torch.ByteTensor([[0, 0],[1,1]])
x.masked_select(indices)
</code></pre>
<p>Slice indexing into matrix:</p>
<pre><code class="lang-auto">x = torch.randn(2, 2)
indices = torch.LongTensor([1, 0])
x.index_select(0, indices)
</code></pre>
<p>Indexing into a 3-tensor:</p>
<pre><code class="lang-auto">x = torch.randn(1, 4, 2)
x[:,:,1]
</code></pre>
<p>Batched indexing into a matrix:</p>
<pre><code class="lang-auto">x = torch.randn(2, 2)
indices = torch.LongTensor([[0, 0],[1,1]])
[x[i] for i in indices]
</code></pre>

gwg: <p>Are you sure these snippets are correct? Just comparing to <a href="https://www.tensorflow.org/api_docs/python/tf/gather_nd" rel="nofollow noopener">TensorFlow’s example</a>:</p>
<pre><code class="lang-auto">indices = torch.ByteTensor([[0, 0], [1, 1]])
params  = torch.Tensor([[1, 2], [3, 4]])
params.masked_select(indices)
</code></pre>
<p>Tensorflow documentation says the output should be <code>tensor([ 1., 4.])</code>, but your code gives <code>tensor([ 3., 4.])</code></p>

richard: <p>That one should be</p>
<pre><code class="lang-auto">indices = torch.tensor([[ 0,  1],[ 0,  1]])
params  = torch.tensor([[1, 2], [3, 4]])
params[indices.tolist()]
</code></pre>

xinwei_he: <p>it should be</p>
<pre><code class="lang-auto">rows_list = [0, 1]
col_list = [0, 1]
params[[rows_list , col_list]]
</code></pre>

palimboa: <p>Are you not missing something in the <code>Batched indexing into a matrix</code> block at the end? If you do it that way you have to loop over all indices, for the dim=0 in your case. My question would be, is there a fast way in pytorch to do the <code>gather_nd</code> where I have a 3D-matrix that stores all the indices and a 3D-matrix that has all the values and I would like to create a new 3D-matrix where each value is mapped to the indices from the 3D-matrix.</p>

richard: <p>Could you show me an example of what you were thinking?<br>
<code>torch.gather</code> can do something like this (<a href="https://pytorch.org/docs/master/torch.html?highlight=gather#torch.gather">https://pytorch.org/docs/master/torch.html?highlight=gather#torch.gather</a>) but it might not be exactly what you’re looking for.</p>

palimboa: <p>I have found the solution myself. Thank you anyway!</p>

Sau_Tak: <p>Hi, how did you solve the problem? Would you mind showing your code? Thx.<img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=6" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:"></p>

erwin: <p>@ palimboa Can you share width me  how do you sovle it ? Thank you</p>

erwin: <p>hello , do you solve it? I also meets the problem</p>

Sau_Tak: <p>emmmmm…, it might work, i only figured out such a naive solution:<br>
the indices must be a 2d tensor and indices.size(1)==len(params.size()).</p>
<pre><code class="lang-auto">def gather_nd(params, indices, name=None):
    '''
    the input indices must be a 2d tensor in the form of [[a,b,..,c],...], 
    which represents the location of the elements.
    '''
    
    indices = indices.t().long()
    ndim = indices.size(0)
    idx = torch.zeros_like(indices[0]).long()
    m = 1
    
    for i in range(ndim)[::-1]:
        idx += indices[i] * m 
        m *= params.size(i)
    
    return torch.take(params, idx)
</code></pre>

erwin: <p>Great! Thank you. I think it’s useful.</p>

Cogito2012: <p>Could you provide a more general version to fulfill tf.gather_nd, not only limited to 2d tensor?</p>
<p>For example, if I want to gather 3d tensor with size CxHxW by giving the indices with size Nx2, where the indices values correspond to the coordinates in HxW grid, how can I get the results with size CxN?</p>
<p>Actually in tensorflow, this can be easily implemented by tf.gather_nd…</p>

ptrblck: <p>I think you can use advanced indexing for this:</p>
<pre><code class="lang-python">C, H, W = 3, 4, 4
x = torch.arange(C*H*W).view(C, H, W)
print(x)
idx = torch.tensor([[0, 0],
                    [1, 1],
                    [2, 2],
                    [3, 3]])

print(x[list((torch.arange(x.size(0)), *idx.chunk(2, 1)))])
</code></pre>

Cogito2012: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thanks! I have found your <a href="https://discuss.pytorch.org/t/extract-batchwise-multiple-coordinates-with-pytorch-functions-similar-to-tf-gather-nd/28097/5?">previous answer</a> which solved my problem.</p>
<p>The following codes are used for validation:</p>
<pre><code class="lang-auto">import torch

batch_size = 2
c, h, w = 256, 38, 65
nb_points = 784
nb_regions = 128

img_feat = torch.randn(batch_size, c, h, w).cuda()
x = torch.empty(batch_size, nb_regions, nb_points, dtype=torch.long).random_(h).cuda()
y = torch.empty(batch_size, nb_regions, nb_points, dtype=torch.long).random_(w).cuda()

# method 1
result_1 = img_feat[torch.arange(batch_size)[:, None], :, x.view(batch_size, -1), y.view(batch_size, -1)]
result_1 = result_1.view(batch_size, nb_regions, nb_points, -1)

# method 2
result_2 = img_feat.new(batch_size, nb_regions, nb_points, img_feat.size(1)).zero_()
for i in range(batch_size):
    for j in range(x.shape[1]):
        for k in range(x.shape[2]):
            result_2[i, j, k] = img_feat[i, :, x[i, j, k].long(), y[i, j, k].long()]

print((result_1 == result_2).all())
</code></pre>

ksb: <p>Thank you for this solution. How can I do this if I have images in batches, NXCXHXW and index tensor are also in batches NXH*WX2. So that My output is NXCXHXW. It would be really great.</p>

ptrblck: <p>Could you post a minimal example, how the index tensor should be used to index which values in your input?</p>

ksb: <p>I have been achieving this using map following is a code and output</p>
<pre><code class="lang-auto">N, C, H, W = 2 , 3 , 2, 2
img = torch.arange(N*C*H*W).view(N,C, H, W)
idx = torch.randint(0,2,(N,H*W,2))

maskit = lambda x,idx: x[list((torch.arange(x.size(0)), *idx.chunk(2, 1)))]

masked = torch.stack([*map(lambda x:maskit(x[0],x[1]),zip(img,idx))])
final = masked.expand(1,*masked.shape).permute(1,3,2,0).view(*img.shape)
</code></pre>
<p>OUTPUT</p>
<pre><code class="lang-auto">print(img)
**out :** tensor([[[[ 0,  1],
          [ 2,  3]],

         [[ 4,  5],
          [ 6,  7]],

         [[ 8,  9],
          [10, 11]]],


        [[[12, 13],
          [14, 15]],

         [[16, 17],
          [18, 19]],

         [[20, 21],
          [22, 23]]]])

print(idx)
**out :** tensor([[[0, 1],
         [1, 0],
         [1, 0],
         [1, 1]],

        [[1, 1],
         [0, 1],
         [1, 0],
         [0, 0]]])
print(masked)
**out :** tensor([[[ 1,  5,  9],
         [ 2,  6, 10],
         [ 2,  6, 10],
         [ 3,  7, 11]],

        [[15, 19, 23],
         [13, 17, 21],
         [14, 18, 22],
         [12, 16, 20]]])

print(final)
**out :** tensor([[[[ 1,  2],
          [ 2,  3]],

         [[ 5,  6],
          [ 6,  7]],

         [[ 9, 10],
          [10, 11]]],


        [[[15, 13],
          [14, 12]],

         [[19, 17],
          [18, 16]],

         [[23, 21],
          [22, 20]]]])

print(final.shape)

**out :** torch.Size([2, 3, 2, 2])

</code></pre>
<p>Can’t we do this without using map.</p>