Apply a coordinate mask on a grid

username: MehdiZouitine
cooked: <p>Hey pytorcher !</p>
<p>As part of a reinforcement learning problem, I need to save gpu memory.</p>
<p>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.</p>
<p>To save this memory I created another matrix format.</p>
<p>I have a mask M shape [Batch,N,2]:</p>
<ul>
<li>B : Batch size</li>
<li>N : Number of agents</li>
<li>2 : x and y coordinates of the agents.</li>
</ul>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/9/d95e9df8f8223d034cbaa6d22bfaa0f726e472fe.png" data-download-href="https://discuss.pytorch.org/uploads/default/d95e9df8f8223d034cbaa6d22bfaa0f726e472fe" title="mask_corrected"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d95e9df8f8223d034cbaa6d22bfaa0f726e472fe_2_690x402.png" alt="mask_corrected" data-base62-sha1="v0WkhW67vQ1JpdXMseExarRYXBc" width="690" height="402" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d95e9df8f8223d034cbaa6d22bfaa0f726e472fe_2_690x402.png, https://discuss.pytorch.org/uploads/default/original/3X/d/9/d95e9df8f8223d034cbaa6d22bfaa0f726e472fe.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/9/d95e9df8f8223d034cbaa6d22bfaa0f726e472fe.png 2x" data-dominant-color="E6D8B2"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">mask_corrected</span><span class="informations">866×505 27.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I also have a shape action grid [Batch,C,H,W]:</p>
<ul>
<li>B : Batch size</li>
<li>C: Number of actions</li>
<li>H : Height of the grid</li>
<li>W: Grid width</li>
</ul>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/6/169b38f23eb806d32b6223bfed0fa749e24a87b9.png" data-download-href="https://discuss.pytorch.org/uploads/default/169b38f23eb806d32b6223bfed0fa749e24a87b9" title="Agent_grid"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/169b38f23eb806d32b6223bfed0fa749e24a87b9_2_477x499.png" alt="Agent_grid" data-base62-sha1="3dZ49bV1ZoODBa3dLHvxZRFUOZX" width="477" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/169b38f23eb806d32b6223bfed0fa749e24a87b9_2_477x499.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/169b38f23eb806d32b6223bfed0fa749e24a87b9_2_715x748.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/6/169b38f23eb806d32b6223bfed0fa749e24a87b9.png 2x" data-dominant-color="ECC1C0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Agent_grid</span><span class="informations">932×975 44.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>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].</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/2/d2f578aca8cd01b575b6d680aeaa7147014c98bf.png" data-download-href="https://discuss.pytorch.org/uploads/default/d2f578aca8cd01b575b6d680aeaa7147014c98bf" title="res_corrected (1)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d2f578aca8cd01b575b6d680aeaa7147014c98bf_2_690x372.png" alt="res_corrected (1)" data-base62-sha1="u6ebHMLqjgTKDHOflZzbXNSH0DB" width="690" height="372" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d2f578aca8cd01b575b6d680aeaa7147014c98bf_2_690x372.png, https://discuss.pytorch.org/uploads/default/original/3X/d/2/d2f578aca8cd01b575b6d680aeaa7147014c98bf.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/2/d2f578aca8cd01b575b6d680aeaa7147014c98bf.png 2x" data-dominant-color="BED2BC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">res_corrected (1)</span><span class="informations">877×474 29.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Here is an example of what I would like to get:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; # Example of dimension
&gt;&gt;&gt; BATCH_SIZE = 2
&gt;&gt;&gt; N_AGENT = 3
&gt;&gt;&gt; NB_ACTION = 2
&gt;&gt;&gt; H_GRID = 3
&gt;&gt;&gt; W_GRID = 3

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

&gt;&gt;&gt; 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]]]])

&gt;&gt;&gt; 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 

&gt;&gt;&gt; output = apply_mask_on_grid(action_grid_batch,batch_mask_agent_position)

&gt;&gt;&gt; output.size()
torch.Size([2, 3, 2]) # [Batch_size, N_AGENT, N_ACTION]
&gt;&gt;&gt; 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
</code></pre>
<p>Thanking you in advance !</p>

username: MehdiZouitine
cooked: <p>Without for loop ! Thanks you !</p>

username: ptrblck
cooked: <p>I’m sure there is a better way of indexing, but this could work for now:</p>
<pre><code class="lang-python">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)
&gt; 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]]])
</code></pre>

username: MehdiZouitine
cooked: <p>Thanks a lot, it work really fast !</p>