对 PyTorch 张量的中间尺寸使用省略号

Question

Suppose I have a torch.Tensor t of shape (8, 3, 32, 32) . I want to index along the first and last 2 dimensions only.

In my usecase, t is a batch of 8 images, of which I want to modify a patch. Suppose the patch is given by indices idx_last = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]]) . I also have idx1 = torch.arange(4) : I want the patch for the first 4 images.

The following does not work:

t[idx1, ..., idx_last]

Is there any way to do this?

Answer 1

I found one workaround, although it may not be the most efficient. In the case where idx1 is 1-dimensional (datapoint selection), and idx_last is multidimensional, the following gets the wanted result:

t[(idx1, ...) + tuple(idx_last.T)]

Better solutions are definitely welcome.