I have a problem that I try to solve, and I have real problems coming up with a solution so far. In tensorflow, one could use tf.gather_nd for it:

Let’s image I have a batch tensor img_feat of size (16, 256, 16, 16). In addition, I have an x and y tensor, each of size (16, 150), so that I have 150 points, each with an x and a y coordinate.

What I want to do is to extract for each batch (16), each of the features (256) for each point (150) which has a x and y coordinate in the 16x16 images of img_feats More precise, this is exactly what I want to have:

feats = torch.stack([torch.stack([img_feat[i,:,x[i,j].long(), y[i,j].long()] for j in range(x.shape[1])]) for i in range(batch_size)])

My problem is, that I want to implement it with torch functions (e.g. gather, select_index, etc.) so that it can be parallelised when using GPUs.

Does anyone can come up with a solution? It would be so awesome!

thanks a lot for your help! I think you are nearly there, however this is not yet what I need:

With your method, points has the shape of (batch_size, c, nb_points, nb_points)

What I want is a resulting tensor of shape (batch_size, nb_points, c). In other words, For each x and y coordinate, I want to extract the whole c features for that coordinate. As I have 150 points per batch, it should result in the above mentioned size.

I attach again my simple list-style code which does the desired behaviour:

feats = torch.stack([torch.stack([img_feat[i,:,x[i,j].long(), y[i,j].long()] for j in range(x.shape[1])]) for i in range(batch_size)])

Thanks so much! I am new to pytorch, and you guys help me a lot!

This way, it should collect only 1 (hopefully the correct ;)) value for each of the nb_points * w extracted values from the first gather in the second gather step.