# Apply a coordinate mask on a grid

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.

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

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].

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]]]])

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.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
``````

2 Likes

Without for loop ! Thanks you !

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

Thanks a lot, it work really fast !

1 Like