How to implement tf.gather_nd with argument batch_dims in Pytorch?

username: ZhengdiYu
cooked: <p>I have been doing a project on image matching, so I need to find correspondences between 2 images. To get descriptors I will need a <strong>interpolate</strong> function. However, when I read about an equivalent function which is done in Tensorflow, I still don’t get how to implement <strong>tf.gather_nd(parmas, indices, barch_dims)</strong> in Pytorch. Especially when there is a argument: batch_dims. I have gone through stackoverflow and there is no perfect equivalence yet.</p>
<p>Here is the interpolate function in Tensorflow and I am trying to implement this in Pytorch:</p>
<p><strong>inputs</strong> is a dense feature map[i] from a for loop of batch size, which means it is 3D[H, W, C](in Pytorch is [C, H, W])<br>
<strong>pos</strong> is a set of random point coordinate shapes like [[I, j], [I, j],…,[I, j]], so it is 2D when it goes in interpolate(in Pytorch is [[i, i,…, i], [j, j,…, j]])</p>
<p><strong>So when doing tf.gather_nd, so the parmas and indices arguments for gather_nd are 4D and 3D</strong></p>
<p>and it then expands both of their dimensions when they get into this function<br>
An example of using it :</p>
<pre><code class="lang-auto">pos = tf.ones((12, 2))  ## stands for a set of coordinates [[i, i,..., i], [j, j,..., j]]

inputs = tf.ones((4, 4, 128))  ## stands for [H, W, C] of dense feature map

outputs = interpolate(pos, inputs, batched=False)

print(outputs.get_shape())  #

&gt;&gt;&gt;(12, 128)
</code></pre>
<p>interpolate function(tensorflow version)</p>
<pre><code class="lang-auto">def interpolate(pos, inputs, nd=True):

    pos = tf.expand_dims(pos, 0)
    inputs = tf.expand_dims(inputs, 0)

    h = tf.shape(inputs)[1]
    w = tf.shape(inputs)[2]

    i = pos[:, :, 0]
    j = pos[:, :, 1]

    i_top_left = tf.clip_by_value(tf.cast(tf.math.floor(i), tf.int32), 0, h - 1)
    j_top_left = tf.clip_by_value(tf.cast(tf.math.floor(j), tf.int32), 0, w - 1)

    i_top_right = tf.clip_by_value(tf.cast(tf.math.floor(i), tf.int32), 0, h - 1)
    j_top_right = tf.clip_by_value(tf.cast(tf.math.ceil(j), tf.int32), 0, w - 1)

    i_bottom_left = tf.clip_by_value(tf.cast(tf.math.ceil(i), tf.int32), 0, h - 1)
    j_bottom_left = tf.clip_by_value(tf.cast(tf.math.floor(j), tf.int32), 0, w - 1)

    i_bottom_right = tf.clip_by_value(tf.cast(tf.math.ceil(i), tf.int32), 0, h - 1)
    j_bottom_right = tf.clip_by_value(tf.cast(tf.math.ceil(j), tf.int32), 0, w - 1)

    dist_i_top_left = i - tf.cast(i_top_left, tf.float32)
    dist_j_top_left = j - tf.cast(j_top_left, tf.float32)
    w_top_left = (1 - dist_i_top_left) * (1 - dist_j_top_left)
    w_top_right = (1 - dist_i_top_left) * dist_j_top_left
    w_bottom_left = dist_i_top_left * (1 - dist_j_top_left)
    w_bottom_right = dist_i_top_left * dist_j_top_left

    if nd:
        w_top_left = w_top_left[..., None]
        w_top_right = w_top_right[..., None]
        w_bottom_left = w_bottom_left[..., None]
        w_bottom_right = w_bottom_right[..., None]

    interpolated_val = (
        w_top_left * tf.gather_nd(inputs, tf.stack([i_top_left, j_top_left], axis=-1), batch_dims=1) +
        w_top_right * tf.gather_nd(inputs, tf.stack([i_top_right, j_top_right], axis=-1), batch_dims=1) +
        w_bottom_left * tf.gather_nd(inputs, tf.stack([i_bottom_left, j_bottom_left], axis=-1), batch_dims=1) +
        w_bottom_right * tf.gather_nd(inputs, tf.stack([i_bottom_right, j_bottom_right], axis=-1), batch_dims=1)
    )

    interpolated_val = tf.squeeze(interpolated_val, axis=0)
    return interpolated_val
</code></pre>

username: ptrblck
cooked: <p><a href="https://discuss.pytorch.org/t/how-to-do-the-tf-gather-nd-in-pytorch/6445/26">This post</a> might provide you a matching implementation for TensorFlow’s <code>gather_nd</code>.<br>
Let me know, if that works for you.</p>

username: ZhengdiYu
cooked: <p>Thank you for your reply! But it doesn’t seem to be compatible with my task</p>
<p>There is an example of using tf.gather_nd in the context</p>
<pre><code class="lang-auto">inputs shape: (1, 120, 160, 128)  ## [N, H, W, C], which stands for dense feature map
indices shape: (1, 512, 2)  ##[N, num_points, channel], which stands for a set of points [ [[iiiii], [jjjjj]] ]
output = tf.gather_nd(inputs, indices, batch_dims = 1)
print(output.get_shape())
&gt;&gt;&gt; (1, 512, 128)
</code></pre>
<p>And just now I find out that in the above interpolate function, we could just simply delete the useless dimension expanding squeezing part at the beginning and the end to get it in a lower dimension case, which can also keep the result meaningful :</p>
<pre><code class="lang-auto">inputs shape: (120, 160, 128)
indices shape: (512, 2)
tf.gather_nd:  (512, 128)
</code></pre>
<p>Is there equivalence in Pytorch to get the equivalent result like this ? :</p>
<pre><code class="lang-auto">inputs :(1, 128, 120, 160)
indices:(1, 2, 512)
gather_nd_torch(inputs, indices)
&gt;&gt;&gt;output:(1, 128, 512)
</code></pre>
<p>or this:</p>
<pre><code class="lang-auto">inputs :(128, 120, 160)
indices:(2, 512)
gather_nd_torch(inputs, indices)
&gt;&gt;&gt;output:(128, 512)
</code></pre>