Pytorch实现将label变成one hot编码的两种方式

由于Pytorch不像TensorFlow有谷歌巨头做维护，很多功能并没有很高级的封装，比如说没有tf.one_hot函数。

本篇介绍将一个mini batch的label向量变成形状为[batch size, class numbers]的one hot编码的两种方法，涉及到

• tensor.scatter_
• tensor.index_select

前言

本文将针对全连接网络和全卷积网络输出的形式不同，将one hot编码分两种情况。

• 第一种针对网络输出是二维，即全连接层的输出形式， [Batchsize， Num_class]
• 第二种针对输出是四维特征图，即分割网络的输出形式，[Batchsize， Num_class， H，W]

先将第一种情况

使用scatter_获得one hot 编码

我相信在CSDN上找这个函数用法的人都是看不懂官方介绍的，所以我不会像其他地方那样，搬官方教程，我也是琢磨了很久才看懂这个函数，但函数声明还是要看看的。

tensor.scatter_(dim, index, src) 

• dim : 指定了覆盖数据是从哪个轴作为依据。后面再详细解释。值的范围是从0到 sum(tensor.shape)-1
• index ： 告诉函数要将src中对应的值放到tensor的哪个位置。index的shape要和src一致，或者src可以通过广播机制实现shape一致。
• src : 保存了想用来覆盖tensor的值

我们先看一个例子，例子从别的博客copy过来，但我会做更加详细的介绍。觉得讲得好请留言作为鼓励。

>>> x = torch.rand(2, 5)
>>> x

 0.4319  0.6500  0.4080  0.8760  0.2355
 0.2609  0.4711  0.8486  0.8573  0.1029
[torch.FloatTensor of size 2x5]

>>> torch.zeros(3, 5).scatter_(0, torch.LongTensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]), x)

 0.4319  0.4711  0.8486  0.8760  0.2355
 0.0000  0.6500  0.0000  0.8573  0.0000
 0.2609  0.0000  0.4080  0.0000  0.1029
[torch.FloatTensor of size 3x5]

注意到dim为0，代表以第一个维度作为依托。index是一个二维数组。

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

那么我们要覆盖tensor的位置有10个，分别为

[0,0];[1,1];[2,2];[0,3];[0,4]
[2,0];[0,1];[0,2];[1,3];[2,4]

dim指定了index我们要将index的值作为哪一个轴的值。其他轴就是按照0到max shape -1变化罢了。比如说dim为0，那么index的值都作为坐标的第一个位置的值，另一个位置从0到4变换。

你们可以验证下，是不是这10个位置被覆盖了。10个位置的第一个轴是index的数字，第二个数字是index中的列数，从0到4。

要覆盖的位置有了，那么用什么值覆盖呢？别忘了我们的index的维度和src是一样的。index中选择什么位置的坐标，就对应用src对应的位置的值代替。

比如说要代替tensor中[0,0]的值，index中[0,0]就是第0行第0列对应的位置，那我们用src第0行第0列的值代替tensor的值。大家可以去验证一下。

我们看看下面的的情况，如果dim为1呢。

>>> z = torch.zeros(2, 4).scatter_(1, torch.LongTensor([[2], [3]]), 1.23)
>>> z

先分析一下

dim为1，那么index的值都作为坐标的第2个位置的值，第一个位置的值应该从0到1变化。

所以要被代替的位置有

[0,2];[1,3]

而[0,2]的位置要填入的值为1.23，[1,3]要填入的值为1.23。（广播机制将1.23这个标量扩展到了shape为（2,1））

好的，函数用法知道了。我们现在看看如何用该函数将label编码为one hot编码。

首先设想一个batch size为8的label。有10类，所以label中的数字应该是从0到9的。

import torch as t
import numpy as np

batch_size = 8
class_num = 10
label = np.random.randint(0,class_num,size=(batch_size,1))
label = t.LongTensor(label)

我们就获得了一个label，shape是(8,1)，必须是2维。如果是(8,)下面的内容会报错的。

y_one_hot = t.zeros(batch_size,class_num).scatter_(1,label,1)
print(y_one_hot)

'''
tensor([[0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],
        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.]])
'''

搞定。下面我们看下面一种方法。

使用tensor.index_select获得one hot编码

还是先看下index_select的用法。

tensor.index_select( dim, index, out=None)

• dim: 指定按什么维度取tensor中的向量
• index: 是一个一维的张量。描述了按照dim维度取出tensor对应的index值的向量。

我们不看例子了，直接看方法，以此为例。

ones = torch.sparse.torch.eye(class_num)
return ones.index_select(0,label)

这里的label是一维的向量，不是二维的。因为index制定了必须是一维的

先生成一个单位矩阵，尺寸是[class_num, class_num]。

dim为0，以为这按照行来取tensor的向量。具体取哪一行呢，就是label中的值了。

这时我们应该也明白为啥这两行代码能实现one hot编码了吧。

如果label是[ 1,3,0]，有四类。那我们得到就是

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

第二种针对分割网络的one_hot编码

对于分割类任务，网络的GT肯定是二维数组，而不是像分类任务那样的一维数组了。而对于分割任务，我们将其视作很多个像素值的分类任务，将ground truth 直接 reshape为向量形式，然后用上面的方法转为one hot编码，然后再reshape回来。核心是不变的。

下面举个例子。

import torch
import numpy as np

gt = np.random.randint(0,5, size=[15,15])  #先生成一个15*15的label，值在5以内，意思是5类分割任务
gt = torch.LongTensor(gt)

def get_one_hot(label, N):
    size = list(label.size())
    label = label.view(-1)   # reshape 为向量
    ones = torch.sparse.torch.eye(N)
    ones = ones.index_select(0, label)   # 用上面的办法转为换one hot
    size.append(N)  # 把类别输目添到size的尾后，准备reshape回原来的尺寸
    return ones.view(*size)


gt_one_hot = get_one_hot(gt, 5)
print(gt_one_hot)
print(gt_one_hot.shape)

print(gt_one_hot.argmax(-1) == gt)  # 判断one hot 转换方式是否正确，全是1就是正确的

另外注意，在Pytorch中，如果要和网络输出的特征图一起计算loss，还要把上面输出的one hot编码的最后一个维度使用permute转到通道维度上。

总结

以上为个人经验，希望能给大家一个参考，也希望大家多多支持。