Apply a coordinate mask on a grid

1

Hey pytorcher !

As part of a reinforcement learning problem, I need to save gpu memory.

The memory size of the grid is quadratic (in number of pixels). However, only the pixels where the agents are located are important. By using “sparse” representations of these grids we save a lot of memory.

To save this memory I created another matrix format.

I have a mask M shape [Batch,N,2]:

  • B : Batch size
  • N : Number of agents
  • 2 : x and y coordinates of the agents.

mask_corrected
mask_corrected866×505 27.9 KB

I also have a shape action grid [Batch,C,H,W]:

  • B : Batch size
  • C: Number of actions
  • H : Height of the grid
  • W: Grid width

Agent_grid
Agent_grid932×975 44.3 KB

The goal is to mask this action grid to obtain only the actions at the coordinate of the mask to obtain a masked grid of shape: [B,N,C].

res_corrected (1)
res_corrected (1)877×474 29.6 KB

Here is an example of what I would like to get:

>>> # Example of dimension
>>> BATCH_SIZE = 2
>>> N_AGENT = 3
>>> NB_ACTION = 2
>>> H_GRID = 3
>>> W_GRID = 3

>>> action_grid_batch.size() # [BATCH_SIZE, NB_ACTION, H_GRID, W_GRID]
torch.Size([2, 2, 3, 3])

>>> action_grid_batch
tensor([[[[0.4000, 0.5000, 0.7000], # Probability  of the action 1 on the action_grid 1 in the batch
          [0.3000, 0.2000, 0.1000],
          [0.9000, 0.8000, 0.7000]],

         [[0.6000, 0.5000, 0.3000], # Probability  of the action 2 on the action_grid 1 in the batch
          [0.7000, 0.8000, 0.9000],
          [0.1000, 0.2000, 0.3000]]],


        [[[0.3000, 0.2000, 0.1000], # Probability  of the action 1 on the action_grid 2 in the batch
          [0.6000, 0.7000, 0.4000],
          [0.9000, 0.8000, 0.1000]],

         [[0.7000, 0.8000, 0.9000], # Probability  of the action 2 on element 2 in the batch
          [0.4000, 0.3000, 0.6000],
          [0.1000, 0.2000, 0.9000]]]])

>>> batch_mask_agent_position
tensor([[[0, 1], # Position (H, W) of the agent 1 on element 1 in the batch
         [1, 1], # Position (H, W) of the agent 2 on element 1 in the batch
         [2, 0]],# Position (H, W) of the agent 3 on element 1 in the batch 

        [[1, 1], # Position (H, W) of the agent 1 on element 2 in the batch
         [1, 2], # Position (H, W) of the agent 2 on element 2 in the batch
         [2, 2]]]) # Position (H, W) of the agent 3 on element 2 in the batch 

>>> output = apply_mask_on_grid(action_grid_batch,batch_mask_agent_position)

>>> output.size()
torch.Size([2, 3, 2]) # [Batch_size, N_AGENT, N_ACTION]
>>> output
tensor([[[0.5000, 0.5000], # Probability of the action 1 and 2 for the agent position 1 on element 1 in the batch
         [0.2000, 0.8000], # Probability of the action 1 and 2 for the agent position 2 on element 1 in the batch
         [0.9000, 0.1000]], # Probability of the action 1 and 2 for the agent position 3 on element 1 in the batch

        [[0.7000, 0.3000], # Probability of the action 1 and 2 for the agent position 1 on element 2 in the batch
         [0.4000, 0.6000], # Probability of the action 1 and 2 for the agent position 2 on element 2 in the batch
         [0.1000, 0.9000]]]) # Probability of the action 1 and 2 for the agent position 3 on element 2 in the batch

Thanking you in advance !

2 Likes
2

Without for loop ! Thanks you !

3

I’m sure there is a better way of indexing, but this could work for now:

a, b = batch_mask_agent_position.split(1, dim=2)
ret = action_grid_batch[torch.arange(action_grid_batch.size(0)), :, a, b]
ret = ret[torch.arange(ret.size(0)), :, torch.arange(ret.size(2))]
print(ret)
> tensor([[[0.5000, 0.5000],
           [0.2000, 0.8000],
           [0.9000, 0.1000]],

          [[0.7000, 0.3000],
           [0.4000, 0.6000],
           [0.1000, 0.9000]]])
3 Likes
4

Thanks a lot, it work really fast !

1 Like