Pytorch下label的one-hot形式转换方法

进行网络训练时通常需要对label转为one-hot形式，下面给出自己知道的两种方法。

方法一

巧妙使用torch.eye()方法

torch.eye(n, m=None, out=None)

参数:

• n (int ) – 行数
• m (int, optional) – 列数.如果为None,则默认为n
• out (Tensor, optinal) - Output tensor

返回一个二维向量，对角线是1，其它位置都是0。

import torch
def one_hot(x, class_count):
	# 第一构造一个[class_count, class_count]的对角线为1的向量
	# 第二保留label对应的行并返回
	return torch.eye(class_count)[x,:]

例子：

x = [0,2,5,4]
class_count = 8
ont_hot(x, class_count)
# 结果
tensor([[1., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 1., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 1., 0., 0.],
[0., 0., 0., 0., 1., 0., 0., 0.]])

返回的是一个二维向量。

方法二

使用scatter_()方法

scatter_(dim, index, src) 

参数

• dim (int) – 索引的轴
• index (LongTensor) – 要散布的元素的索引。可以为空或src的大小相同， 如果为空，则操作返回标识。如果要用于label -> one-hot，其实这里就是label（因为label的大小就对应了它的位置，注意label是二维大小！）
• src (Tensor) – 指定填充的元素值

例子：

# 第一 先构造[2,4]维度的二维向量
# 第二 在1轴的2，3索引处填值 1.23
>>> z = torch.zeros(2, 4).scatter_(1, torch.tensor([[2], [3]]), 1.23)
>>> z
tensor([[ 0.0000,  0.0000,  1.2300,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  1.2300]])

import torch as t
import numpy as np


# 首先设想一个batch size为8的label，10类
batch_size = 8
class_num = 10
label = np.random.randint(0,class_num,size=(batch_size,1))
# 获得了一个label，shape是(8,1)，必须是2维
label = t.LongTensor(label)


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

参考
https://pytorch.org/docs/stable/tensors.html