Segmentation to one-hot encoding
I have a batch of segmented images of size
seg --&gt; [batch, channels, imsize, imgsize] --&gt; [16, 6, 50, 50]
each scalar in this tensor specifies one of the segmentation classes.
We have 2000 total segmentation classes.
Now the goal is to convert
[16, 6, 50, 50] --&gt; [16, 2000, 50, 50]
where each class is encoded in one hot fashion.
How do I do it with pytorch api?
I can only think of ridiculously inefficient looping construction.
Example
Here we will have only 2 initial channels (instead of 6), 4 labels (instead of 2000), size batch 1 (instead of 16) and 4x4 image instead of 50x50.
0, 0, 1, 1
0, 0, 0, 1
1, 1, 1, 1
1, 1, 1, 1

3, 3, 2, 2
3, 3, 2, 2
3, 3, 2, 2
3, 3, 2, 2

Now this turns into 4 channel output
1, 1, 0, 0
1, 1, 1, 0
0, 0, 0, 0
0, 0, 0, 0

0, 0, 1, 1
0, 0, 0, 1
1, 1, 1, 1
1, 1, 1, 1

1, 1, 0, 0
1, 1, 0, 0
1, 1, 0, 0
1, 1, 0, 0

0, 0, 1, 1
0, 0, 1, 1
0, 0, 1, 1
0, 0, 1, 1

The key observation is that a particular label appears only on a single input channel.

I think you can achieve this without too much trouble. Construct as many masks as there are labels, then stack those masks together, sum on the channel layer and convert to floats:
&gt;&gt;&gt; x
tensor([[[0, 0, 1, 1],
         [0, 0, 0, 1],
         [1, 1, 1, 1],
         [1, 1, 1, 1]],

        [[3, 3, 2, 2],
         [3, 3, 2, 2],
         [3, 3, 2, 2],
         [3, 3, 2, 2]]])

&gt;&gt;&gt; y = torch.stack([x==i for i in range(x.max()+1)], dim=1).sum(dim=2)
tensor([[[1., 1., 0., 0.],
         [1., 1., 1., 0.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.]],

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

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

        [[1., 1., 0., 0.],
         [1., 1., 0., 0.],
         [1., 1., 0., 0.],
         [1., 1., 0., 0.]]])

