How to get a batch of binary masks from segmentation map without using `for` loop

username: qiminchen
cooked: <p>Hi, suppose I have a segmentation map <code>a</code> with dimension of <code>torch.Size([1, 1, 6, 6])</code></p>
<pre><code class="lang-auto">print(a)
tensor([[[[ 0.,  0.,  0.,  0.,  0.,  0.],
          [ 0., 15., 15., 16., 16.,  0.],
          [ 0., 15., 15., 16., 16.,  0.],
          [ 0., 13., 13.,  9.,  9.,  0.],
          [ 0., 13., 13.,  9.,  9.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.]]]])
</code></pre>
<p>How can I get the binary masks for each id without using <code>for</code> loop? The binary masks should have a dimension of <code>torch.Size([1, 4, 6, 6])</code>, currently im doing something like this and the reason I want it without <code>for</code> loop is that the dimension of <code>a</code> might change and there might be more/less classes. Thanks.</p>
<pre><code class="lang-auto">a1 = torch.where(segmentation_a == 15, 1, 0)
a2 = torch.where(segmentation_a == 16, 1, 0)
a3 = torch.where(segmentation_a == 13, 1, 0)
a4 = torch.where(segmentation_a == 9, 1, 0)
b = torch.cat((a1, a2, a3, a4), dim=1)

print(b)
tensor([[[[0, 0, 0, 0, 0, 0],
          [0, 1, 1, 0, 0, 0],
          [0, 1, 1, 0, 0, 0],
          [0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0]]],


        [[[0, 0, 0, 0, 0, 0],
          [0, 0, 0, 1, 1, 0],
          [0, 0, 0, 1, 1, 0],
          [0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0]]],


        [[[0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0],
          [0, 1, 1, 0, 0, 0],
          [0, 1, 1, 0, 0, 0],
          [0, 0, 0, 0, 0, 0]]],


        [[[0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 0, 0],
          [0, 0, 0, 1, 1, 0],
          [0, 0, 0, 1, 1, 0],
          [0, 0, 0, 0, 0, 0]]]])
</code></pre>

username: eqy
cooked: <p>Does <code>torch.nn.functional.one_hot</code> work for you?</p>
<pre><code class="lang-auto">import torch
import time
a = torch.tensor([[[[ 0.,  0.,  0.,  0.,  0.,  0.],
                    [ 0., 15., 15., 16., 16.,  0.],
                    [ 0., 15., 15., 16., 16.,  0.],
                    [ 0., 13., 13.,  9.,  9.,  0.],
                    [ 0., 13., 13.,  9.,  9.,  0.],
                    [ 0.,  0.,  0.,  0.,  0.,  0.]]]])
# using CPU timing so playing fast and loose with warmup/syncs
t1 = time.time()
a1 = torch.where(a == 15, 1, 0)
a2 = torch.where(a == 16, 1, 0)
a3 = torch.where(a == 13, 1, 0)
a4 = torch.where(a == 9, 1, 0)
b = torch.cat((a1, a2, a3, a4), dim=1)
t2 = time.time()
o = torch.nn.functional.one_hot(a.long())
t3 = time.time()
print(torch.all(a1 ==  o[:,:,:,:,15]))
print(torch.all(a2 ==  o[:,:,:,:,16]))
print(torch.all(a3 ==  o[:,:,:,:,13]))
print(torch.all(a4 ==  o[:,:,:,:,9]))
print(f"loop took {t2-t1} one_hot took {t3-t2}")
</code></pre>
<pre><code class="lang-auto">tensor(True)
tensor(True)
tensor(True)
tensor(True)
loop took 0.0006623268127441406 one_hot took 0.0002810955047607422
</code></pre>

username: qiminchen
cooked: <p>Hi <a class="mention" href="/u/eqy">@eqy</a> thanks for the idea, I just found another solution:</p>
<pre><code class="lang-auto">b = torch.where(a == torch.tensor([15, 16, 13, 9]).view(1, 4, 1, 1), 1, 0)
</code></pre>