Pytorch Set value for Tensor with index tensor

user_6034273 0
Viewed 1957

Suppose i have a 2d index Array of shape [B,1,N,2] i.e N points holding indexes on a target tensor of size [B,1,H,W]. what is the best way to assign value to all the indices in the tensor?

for example:

for b in batchsize:
  for i in N: 
    target[b,0,ind[i,0],ind[i,1]] = 1

but not in loop form thanks

pytorch python
2 Answers

If we look at this setup, you have a tensor target shaped (b, 1, h, w) and a tensor containing indices ind, shaped (b, 1, N, 2). You want to assign 1 to the N points per batch given by the two coordinates in ind.

The way I see it you could use torch.scatter_. We will stick with a 3D tensor since axis=1 is unused. Given a 3D tensor and argument value=1 and dim=1, .scatter_ operates on the input tensor as so:

input[i][index[i][j]] = 1

This does not exactly fit your setting since what would wish for is rather

input[i][index1[i][j]][index2[i][j]] = 1

In order to use scatter you could flatten target and unfold the values in ind accordingly.

Let's take an example:

>>> target = torch.zeros(2, 1, 10, 10)
>>> ind = torch.tensor([[[[0, 0], [1, 1], [2, 2]]],
                        [[[1, 2], [3, 4], [7, 8]]]])
tensor([[[[0, 0],
          [1, 1],
          [2, 2]]],


        [[[1, 2],
          [3, 0],
          [4, 2]]]])

We will start by splitting ind into xs and ys coordinates:

>>> x, y = ind[..., 0], ind[..., 1]

Unfold and reshape them:

>>> z = x*target.size(-1) + y
tensor([[[ 0,  4,  8]],

        [[ 5,  9, 14]]])

Flatten target:

>>> t = target.flatten(2)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])

Then scatter the 1s:

>>> t.scatter_(dim=2, index=z, value=1)
tensor([[[1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.]],

        [[0., 0., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0., 1.]]])

And finally reshape back to the original shape:

>>> t.reshape_as(target)
tensor([[[[1., 0., 0.],
          [0., 1., 0.],
          [0., 0., 1.],
          [0., 0., 0.],
          [0., 0., 0.]]],

        [[[0., 0., 0.],
          [0., 0., 1.],
          [0., 0., 0.],
          [1., 0., 0.],
          [0., 0., 1.]]]])

In summary:

>>> x, y = ind[..., 0], ind[..., 1]
>>> z = x*target.size(-1) + y
>>> target.flatten(2).scatter_(dim=2, index=z, value=1).reshape_as(target)

This last line will mutate target.

user_63313690

It's a bit hard to tell what you mean by "a 2D array of shape [B,1,N,2]" (that looks like a 4D array in commonly used notation). But it looks like across all elements in the batch, you want to assign the same values on dimensions 2 and 3 to each element indexed by N. You can use list indexing notation:

# assuming ind is a list of tuples of (dim2 index, dim3 index)

# unzip ind such that each resulting list contains indices along a single dimension
dim2_indices = [item[0] for item in ind]
dim3_indices = [item[1] for item in ind]
target[:,0,dim2,dim3] = 1
user_98317770
Related