分割为 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.]]])