Masking tensor with multiple masks without for loop

melike: <p>Hi all, I have a 5-dimensional tensor of masks and I want to create a 4-D tensor by masking a 4-D input with each of these masks. So far I could come up with a for loop on unbinded masks, but I’d appreciate if someone could suggest a way to eliminate the for loop since the number of masks is high (hundreds). Here is an example,</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch

&gt;&gt;&gt; _ = torch.manual_seed(16)
&gt;&gt;&gt; num_masks = 2
&gt;&gt;&gt; inp = torch.randint(1, 5, (3, 4), dtype=torch.float)
&gt;&gt;&gt; outp = torch.zeros(inp.shape)
&gt;&gt;&gt; masks = torch.randn((num_masks,) + inp.shape) &gt; 0
&gt;&gt;&gt; for m in torch.unbind(masks, dim=0):
&gt;&gt;&gt; 	outp[m] = inp[m]


&gt;&gt;&gt; inp
tensor([[2., 3., 2., 2.],
        [2., 4., 1., 2.],
        [4., 1., 2., 1.]])

&gt;&gt;&gt; masks
tensor([[[False,  True,  True,  True],
         [ True, False, False,  True],
         [ True, False,  True, False]],

        [[False,  True,  True,  True],
         [ True,  True, False, False],
         [False,  True,  True,  True]]])

&gt;&gt;&gt; outp
tensor([[0., 3., 2., 2.],
        [2., 4., 0., 2.],
        [4., 1., 2., 1.]])
</code></pre>

ptrblck: <p>I don’t think there is a proper way to avoid the for loop, since you are sequentially overwriting <code>out</code> in:</p>
<pre><code class="lang-python">outp[m] = inp[m]
</code></pre>
<p>I.e. <code>outp</code> and <code>inp</code> have a shape of <code>[3, 4]</code> while you are creating <code>num_masks</code> with the same shape.<br>
In the first iteration the first mask would thus manipulate <code>outp</code>, in the second iteration this already changed <code>outp</code> tensor would then be again manipulated so you have a sequential dependency.</p>
<p>One potential way would be to “sum” all masks to a single one and perform this operation only once.</p>
<pre><code class="lang-python">out = torch.zeros(inp.shape)
mask = masks.sum(dim=0).bool()
out[mask] = inp[mask]
out == outp
</code></pre>
<p>but it depends on your use case and if that’s really what you are trying to achieve.</p>

melike: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>Summing the masks would not be an appropriate approach to this problem. Each mask represents a connected component (of arbitrary shape) in the image. I use them for a custom max pooling operation. So for each mask, all pixels of that CC are set to their maximum value.</p>
<p>The operation does not have to be sequential, but I could not think of a better way than the for loop, so any non-sequential approaches are especially welcome. Here is an example and I can post its whole code if needed but I’m basically stuck at the for loop mentioned in the 1st post.</p>
<pre><code class="lang-auto">&gt;&gt;&gt; inp
tensor([[[[11., 11., 10., 12.],
          [10., 11., 11., 10.],
          [11., 12., 12., 12.],
          [10., 12., 11., 10.]]]], requires_grad=True)

&gt;&gt;&gt; CCs
tensor([[0., 0., 0., 1.],
        [2., 2., 2., 2.],
        [2., 2., 2., 2.],
        [2., 2., 2., 2.]])

&gt;&gt;&gt; masks
tensor([[[[[ True,  True,  True, False],
           [False, False, False, False],
           [False, False, False, False],
           [False, False, False, False]]]],

        [[[[False, False, False,  True],
           [False, False, False, False],
           [False, False, False, False],
           [False, False, False, False]]]],

        [[[[False, False, False, False],
           [ True,  True,  True,  True],
           [ True,  True,  True,  True],
           [ True,  True,  True,  True]]]]])

&gt;&gt;&gt; output
tensor([[[[11., 11., 11., 12.],
          [12., 12., 12., 12.],
          [12., 12., 12., 12.],
          [12., 12., 12., 12.]]]])
</code></pre>

ptrblck: <p>Unfortunately, I don’t understand how your <code>output</code> tensor is created given the <code>inp</code> and <code>masks</code> (e.g. where is the <code>11.</code> coming from in <code>output[0, 0, 0, 2]</code>).</p>
<p>In any case, using your manual loop approach and the “summed” mask approach would still yield the same result, as you are just applying the masks sequentially and could thus also just create a single mask:</p>
<pre><code class="lang-python"># setup
inp = torch.tensor([[[[11., 11., 10., 12.],
                      [10., 11., 11., 10.],
                      [11., 12., 12., 12.],
                      [10., 12., 11., 10.]]]], requires_grad=True)

masks = torch.tensor([[[[[ True,  True,  True, False],
                         [False, False, False, False],
                         [False, False, False, False],
                         [False, False, False, False]]]],

                      [[[[False, False, False,  True],
                         [False, False, False, False],
                         [False, False, False, False],
                         [False, False, False, False]]]],

                      [[[[False, False, False, False],
                         [ True,  True,  True,  True],
                         [ True,  True,  True,  True],
                         [ True,  True,  True,  True]]]]])

# reference
reference_output = torch.tensor([[[[11., 11., 11., 12.],
                                   [12., 12., 12., 12.],
                                   [12., 12., 12., 12.],
                                   [12., 12., 12., 12.]]]])

# loop
loop_output = torch.zeros(inp.shape)
for m in torch.unbind(masks, dim=0):
 	loop_output[m] = inp[m]

print((loop_output == reference_output).all())
&gt; tensor(False)

# summed
mask = masks.sum(dim=0).bool()
summed_output = torch.zeros(inp.shape)
summed_output[mask] = inp[mask]
print((summed_output==reference_output).all())
&gt; tensor(False)
print((summed_output==loop_output).all())
&gt; tensor(True)
</code></pre>

melike: <blockquote>
<p>Unfortunately, I don’t understand how your <code>output</code> tensor is created given the <code>inp</code> and <code>masks</code> (e.g. where is the <code>11.</code> coming from in <code>output[0, 0, 0, 2]</code>).</p>
</blockquote>
<p>Sorry for the poor explanation. Each value in <code>CCs</code> indicates which connected component that pixel belongs to. So in the above example, 3 pixels belong to 0th CC, 1 pixel belongs to 1st CC and the rest belong to 2nd CC.</p>
<p>For custom pooling, all pixels of a CC are set to their maximum value from the input image. These are <code>11.</code> for the 0th CC, <code>12.</code> for the 1st CC and again <code>12.</code> for the 2nd CC.</p>
<p>The problem with the summed mask approach is that each pixel exists in a connected component. Therefore, <code>masks.sum(dim=0).bool()</code> results in a <code>mask</code> filled only with <code>True</code> values, and applying it to <code>input</code> will return <code>input</code> again. I added the operation, the for loop in <code>forward</code> is causing a bottleneck.</p>
<pre><code class="lang-auto">import torch

class AdvMaxify(torch.autograd.Function):
    
    @staticmethod
    def forward(ctx, inp, CCs):
        maxified = torch.ones(inp.shape, device=inp.device)
        cs = torch.unique(CCs)
        mask_shape = cs.shape + maxified.shape
        masks = CCs.expand(mask_shape) == cs.view(cs.shape + (1, 1, 1, 1)).expand(mask_shape)
        for m in torch.unbind(masks, dim=0): # bottleneck
            m_inp = inp * m
            dim_maxes = torch.amax(m_inp.view(inp.shape[0], inp.shape[1], -1).squeeze(), dim=-1)
            num_dims = m.dim() - dim_maxes.dim()
            new_shape = dim_maxes.shape + (1,) * num_dims if m.shape[0] &gt; 1 else (1,) + dim_maxes.shape + (1,) * (num_dims - 1)
            maxified[m] = dim_maxes.view(new_shape).expand_as(m)[m]
        return maxified
    
    @staticmethod
    def backward(ctx, grad_output):
        return grad_output, None


if __name__ == "__main__":

    inp = torch.tensor([[[[11., 11., 10., 12.],
                          [10., 11., 11., 10.],
                          [11., 12., 12., 12.],
                          [10., 12., 11., 10.]]]], requires_grad=True)

    CCs = torch.tensor([[0., 0., 0., 1.],
                        [2., 2., 2., 2.],
                        [2., 2., 2., 2.],
                        [2., 2., 2., 2.]])

    reference_output = torch.tensor([[[[11., 11., 11., 12.],
	                                   [12., 12., 12., 12.],
	                                   [12., 12., 12., 12.],
	                                   [12., 12., 12., 12.]]]])

    adv_maxify = AdvMaxify.apply
    outp = adv_maxify(inp, CCs)

print((outp == reference_output).all())
&gt; tensor(True)
</code></pre>

tom: <p>If the components are disjoint and you get the numbers in a single integer tensor, you could first use <code>scatter_max</code> from the <a href="https://pytorch-scatter.readthedocs.io/en/latest/index.html" rel="noopener nofollow ugc">PyTorch scatter library</a> and then indexing with the mask to get the values back into the index.<br>
There are faster ways to do this, but at least the one I implemented isn’t open source.</p>
<p>Best regards</p>
<p>Thomas</p>

melike: <p>Hi <a class="mention" href="/u/tom">@tom</a>, yes, the connected components are disjoint.</p>
<blockquote>
<p>and then indexing with the mask to get the values back into the index.</p>
</blockquote>
<p>Could you elaborate this part a bit more please ?</p>
<p><strong>Update:</strong> This is what I could come up from your approach, but it still has the same bottleneck <code>for</code> loop,</p>
<pre><code class="lang-auto">import torch
from torch_scatter import scatter_max

inp = torch.tensor([[[[11., 11., 10., 12.],
                      [10., 11., 11., 10.],
                      [11., 12., 12., 12.],
                      [10., 12., 11., 10.]]]], requires_grad=True)

CCs = torch.tensor([[0, 0, 0, 1],
			        [2, 2, 2, 2],
			        [2, 2, 2, 2],
			        [2, 2, 2, 2]])

reference_output = torch.tensor([[[[11., 11., 11., 12.],
                                   [12., 12., 12., 12.],
                                   [12., 12., 12., 12.],
                                   [12., 12., 12., 12.]]]])


indexes = torch.unique(CCs)
out = inp.new_zeros(indexes.shape)
out, argmax = scatter_max(inp.view(-1), CCs.view(-1), out=out)
maxified = inp.new_ones(inp.shape)

for i in range(len(indexes)):
    mask = indexes[i] == CCs
    maxified[mask.view((1, 1) + mask.shape)] = out[i]

print((maxified == reference_output).all())
&gt; tensor(True)
</code></pre>

melike: <p>I found a way to do it without a <code>for</code> loop, combining the approaches of <a class="mention" href="/u/tom">@tom</a>, <a class="mention" href="/u/ptrblck">@ptrblck</a> and with some support from friends,</p>
<pre><code class="lang-auto">import torch
from torch_scatter import scatter_max

class AdvMaxify(torch.autograd.Function):
    
    @staticmethod
    def forward(ctx, inp, CCs):
        indexes = torch.unique(CCs)                                                     # These are also unique connected component ids. 
        cc_maxes = inp.new_zeros(inp.shape[:2] + indexes.shape)                         # Tensor in [bs, ch, num_CCs] shape to keep each CC's max value. 
        cc_maxes, _ = scatter_max(src=inp.view(inp.shape[:2] + (-1,)),
                                  index=CCs.unsqueeze(1).expand_as(inp).view(inp.shape[:2] + (-1,)),
                                  out=cc_maxes)
        mask_shape = indexes.shape + inp.shape
        masks = CCs.unsqueeze(1).expand(mask_shape) == indexes.view(indexes.shape + (1, 1, 1, 1)).expand(mask_shape)
        return torch.sum(masks * cc_maxes.permute(2, 0, 1).contiguous().view(mask_shape[:3] + (1, 1)).expand(mask_shape), dim=0)


    @staticmethod
    def backward(ctx, grad_output):
        return grad_output, None

if __name__ == "__main__":
    adv_maxify = AdvMaxify.apply
    inp = torch.tensor([[[[ 8.,  5., 10., 12.],
                          [11.,  7., 10., 12.],
                          [10.,  8., 12., 10.],
                          [ 8., 11., 10.,  6.]],
    
                         [[ 7., 10.,  6.,  5.],
                          [ 7.,  7.,  5.,  7.],
                          [12., 10.,  5., 11.],
                          [ 7.,  6.,  7., 12.]],
    
                         [[ 9.,  7.,  5.,  9.],
                          [ 7.,  9., 12.,  9.],
                          [10., 11., 11.,  8.],
                          [10., 10.,  6.,  8.]]],
    
                        [[[10., 11., 11., 10.],
                          [ 8., 10., 12.,  7.],
                          [ 8.,  5.,  7.,  8.],
                          [ 6.,  9.,  6.,  8.]],
    
                         [[ 9., 10.,  7., 11.],
                          [ 8.,  9.,  7., 11.],
                          [ 8., 10.,  6., 11.],
                          [10.,  7.,  7.,  6.]],
    
                         [[12., 10., 10., 11.],
                          [12., 11., 10.,  8.],
                          [12.,  8., 10.,  9.],
                          [12., 10., 10., 10.]]]], requires_grad=True)
    
    cidx = torch.tensor([[[0, 1, 1, 1],
                          [2, 3, 3, 3],
                          [2, 3, 3, 3],
                          [2, 3, 3, 3]],
    
                         [[0, 1, 1, 2],
                          [0, 1, 1, 2],
                          [3, 4, 4, 5],
                          [3, 4, 4, 5]]])
    
    reference_output = torch.tensor([[[[ 8., 12., 12., 12.],
                                       [11., 12., 12., 12.],
                                       [11., 12., 12., 12.],
                                       [11., 12., 12., 12.]],
     
                                      [[ 7., 10., 10., 10.],
                                       [12., 12., 12., 12.],
                                       [12., 12., 12., 12.],
                                       [12., 12., 12., 12.]],
     
                                      [[ 9.,  9.,  9.,  9.],
                                       [10., 12., 12., 12.],
                                       [10., 12., 12., 12.],
                                       [10., 12., 12., 12.]]],
     
                                     [[[10., 12., 12., 10.],
                                       [10., 12., 12., 10.],
                                       [ 8.,  9.,  9.,  8.],
                                       [ 8.,  9.,  9.,  8.]],
     
                                      [[ 9., 10., 10., 11.],
                                       [ 9., 10., 10., 11.],
                                       [10., 10., 10., 11.],
                                       [10., 10., 10., 11.]],
     
                                      [[12., 11., 11., 11.],
                                       [12., 11., 11., 11.],
                                       [12., 10., 10., 10.],
                                       [12., 10., 10., 10.]]]]) 
    
    maxed = adv_maxify(inp, cidx)
    
print(torch.all(reference_output == maxed))
&gt; tensor(True)
</code></pre>

tom: <p>Good job. Congratulations!</p>