Indexing elements from a batch tensor in PyTorch
Say, I have a batch of images in PyTorch. For each image, I also have a pixel location say (x, y). The pixel value can be read using img[x, y] for one image. I am trying to read pixel value for each image in the batch. Please see below the code snippet:
import torch

# create tensors to represent random images in torch format
img_1 = torch.rand(1, 200, 300)
img_2 = torch.rand(1, 200, 300)
img_3 = torch.rand(1, 200, 300)
img_4 = torch.rand(1, 200, 300)

# for each image, x-y value are know, so creating a tuple
img1_xy = (0, 10, 70)
img2_xy = (0, 40, 20)
img3_xy = (0, 30, 50)
img4_xy = (0, 80, 60)

# this is what I am doing right now
imgs = [img_1, img_2, img_3, img_4]
imgs_xy = [img1_xy, img2_xy, img3_xy, img4_xy]
x = [img[xy] for img, xy in zip(imgs, imgs_xy)]
x = torch.as_tensor(x)

My Concerns and Questions

In each image, the pixel location i.e., (x, y) is know. However, I have to create a tuple with one more element, i.e., 0 to make sure that tuple matches the shape of image. Any elegant way?
Instead of using tuple, can't we use tensor and then get the pixel values?
All images can be concatenated to make a batch as img_batch = torch.cat((img_1, img_2, img_3, img_4)). But what about tuple?


You can concatanate the images to form a  (4, 200, 300) shaped stacked tensor. Then, we can index into this with the known (x, y) pairs for each image as follows: we need [0, x1, y1] for first image, [1, x2, y2] for second image, [2, x3, y3] for third image and so on. These can be achieved with &quot;fancy indexing&quot;:
# stacking as you did
&gt;&gt;&gt; stacked_imgs = torch.cat(imgs)
&gt;&gt;&gt; stacked_imgs.shape
(4, 200, 300)

# no need for 0s in front
&gt;&gt;&gt; imgs_xy = [(10, 70), (40, 20), (30, 50), (80, 60)]

# need xs together and ys together: take transpose of `imgs_xy`
&gt;&gt;&gt; inds_x, inds_y = torch.tensor(imgs_xy).T

&gt;&gt;&gt; inds_x
tensor([10, 40, 30, 80])

&gt;&gt;&gt; inds_y
tensor([70, 20, 50, 60])

# now we index into the batch
&gt;&gt;&gt; num_imgs = len(imgs)
&gt;&gt;&gt; result = stacked_imgs[range(num_imgs), inds_x, inds_y]
&gt;&gt;&gt; result
tensor([0.5359, 0.4863, 0.6942, 0.6071])

We can check the result:
&gt;&gt;&gt; torch.tensor([img[0, x, y] for img, (x, y) in zip(imgs, imgs_xy)])

tensor([0.5359, 0.4863, 0.6942, 0.6071])


To answer your questions:
1: Since we stacked the images, that issue is mitigated and we use range(4) to index into each individual image instead.
2: Yes, we indeed turn x, y positions into tensors.
3: We directly index with them after they are separated into tensors.
