分割为one-hot编码

Question

我有一批大小的分割图像

seg --> [batch, channels, imsize, imgsize] --> [16, 6, 50, 50]

该张量中的每个标量指定一个分割类。我们有2000总分割类。

现在的目标是转换[16, 6, 50, 50] --> [16, 2000, 50, 50]，其中每个类都以一种热门的方式进行编码。

如何使用pytorch api完成此操作？我只能想到非常低效的循环构造。

示例

这里我们将只有2个初始通道(而不是6个)，4个标签(而不是2000个)，大小为1号(而不是16个)和4x4图像而不是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

现在，这变成了4通道输出

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

关键的观察结果是，特定的标签只出现在单个输入通道上。

Answer 1

我认为你可以毫不费力地做到这一点。构建与标签数量一样多的蒙版，然后将这些蒙版堆叠在一起，在通道层上求和并转换为浮点数：

>>> 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]]])

>>> 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.]]])