Multidimensional filling in a N dimensional tensor

martinferianc: <p>Hi,</p>
<p>Is there a way, where it would be possible to use a vectorizing operation if we would like to fill a portion of the tensor, given different indices for each sample in a batch <code>B</code>? For example:</p>
<pre data-code-wrap="python3"><code class="lang-plaintext">import torch
x = torch.randn(3, 3, 224, 224)
x_shape = x.shape
B, C, H, W = x_shape
cut_width = int(W * 0.5)
cut_height = int(H * 0.5)

cut_x = torch.randint(0, W, (B,))
cut_y = torch.randint(0, H, (B,))

bbx1 = torch.clamp(cut_x - cut_width // 2, 0, W)
bby1 = torch.clamp(cut_y - cut_height // 2, 0, H)
bbx2 = torch.clamp(cut_x + cut_width // 2, 0, W)
bby2 = torch.clamp(cut_y + cut_height // 2, 0, H)

mask = torch.zeros(x_shape)
         
# Repalce the for loop with a vectorized implementation
for i in range(B):
    mask[i, :, bby1[i]:bby2[i], bbx1[i]:bbx2[i]] = 1.0
</code></pre>
<p>I did come across e.g. <a href="https://pytorch.org/vision/stable/generated/torchvision.transforms.functional.erase.html" class="inline-onebox" rel="noopener nofollow ugc">erase — Torchvision main documentation</a> but it does not work over batched inputs.</p>

KFrank: <p>Hi Martin!</p>
<aside class="quote no-group quote-modified" data-username="martinferianc" data-post="1" data-topic="171048" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/martinferianc/48/18507_2.png" class="avatar"> martinferianc:</div>
<blockquote>
<p>Is there a way, where it would be possible to use a vectorizing operation if we would like to fill a portion of the tensor, given different indices for each sample in a batch <code>B</code>?</p>
<pre><code class="lang-auto"># Repalce the for loop with a vectorized implementation
for i in range(B):
    mask[i, :, bby1[i]:bby2[i], bbx1[i]:bbx2[i]] = 1.0
</code></pre>
</blockquote>
</aside>
<p>In a reply to a similar question, I showed how to create a “zero-mask” from<br>
a batch of bounding boxes.  You should be able to do the same, but convert<br>
the “zero-mask” to a “one-mask” with <code>logical_not()</code>.</p>
<p>Here is the relevant post:</p>
<aside class="quote quote-modified" data-post="2" data-topic="166111">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/batch-indexing-to-avoid-the-loop/166111/2">Batch indexing to avoid the loop</a> 
  </div>
  <blockquote>
    Hi Hwang! 

Yes, you can use your bounding boxes, B, to build a mask tensor that 
you then multiply onto img to zero out the values in the bounding boxes: 
&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.13.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; nBatch = 256
&gt;&gt;&gt; nChannels = 3
&gt;&gt;&gt; h = 32
&gt;&gt;&gt; w = h
&gt;&gt;&gt;
&gt;&gt;&gt; img = torch.randn (nBatch, nChannels, h, w)
&gt;&gt;&gt; B = torch.rand (nBatch, 4) * h
&gt;&gt;&gt; B = B.type (torch.int32)
&gt;&gt;&gt;
&gt;&gt;&gt; imgB = img.clone()
&gt;&gt;&gt;
&gt;&gt;&gt; # loop version
&gt;&gt;&gt; for  i in range (nBatch…
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>

martinferianc: <p>Thank you very much for your post if I “hack” it like this, this works:</p>
<pre><code class="lang-auto">y_mask = torch.logical_or(torch.arange(H).unsqueeze(
            0) &lt; bby1[:, None], torch.arange(H).unsqueeze(0) &gt;= bby2[:, None])
 x_mask = torch.logical_or(torch.arange(W).unsqueeze(
            0) &lt; bbx1[:, None], torch.arange(W).unsqueeze(0) &gt;= bbx2[:, None])
 ones_mask = ~torch.logical_or(y_mask.unsqueeze(2), x_mask.unsqueeze(1)).unsqueeze(1)
 mask = torch.zeros(x_shape)
 mask+=ones_mask
</code></pre>
<p>Is there potentially a more elegant solution? <a class="mention" href="/u/ptrblck">@ptrblck</a> Thank you again!</p>

ptrblck: <p>I think your approach sticks to <a class="mention" href="/u/kfrank">@KFrank</a>’s suggestion and looks alright.<br>
You could remove the last step and directly transform <code>ones_mask</code> to a <code>FloatTensor</code> as it seems to be your target:</p>
<pre><code class="lang-python">ones_mask = ~torch.logical_or(y_mask.unsqueeze(2), x_mask.unsqueeze(1)).unsqueeze(1)
mask = ones_mask.float()
</code></pre>