PyTorch torch.nn.functional.one_hot用法解析

1.用法

在PyTorch中，我们可以使用torch.nn.functional.one_hot函数来实现One-Hot编码。下面是一个简单的例子：

import torch
import torch.nn.functional as F
# 假设我们有一个包含类别标签的张量
labels = torch.tensor([0, 2, 1, 0, 2])
# 使用torch.nn.functional.one_hot进行One-Hot编码
one_hot = F.one_hot(labels, num_classes=3)
print(one_hot)

重点在于下面的两点：

2.one_hot的输入需要是非负整数张量（小数和负数都不行）

3.经过one_hot处理后张量维度的变化：

假设输入的张量维度是n，那么输出张量维度就是n+1，而且多的这一维度是加在了最后一维。例如，输入张量是1维的，经过one_hot处理后就变成了2维的。关于这最后一维具体是多少，又有两种情况：

import torch
import torch.nn.functional as F

a = torch.tensor([[1,1,2]])
b = F.one_hot(a)
c = F.one_hot(a,4)
print(b.shape,c.shape)

A.以上面代码为例，如果不指定num_classes，pytorch默认将a中最大值加1作为标签类别最大数，此时最后一维就等于该最大值。例如，a中最大值是2，标签类别最大数就是2+1=3，那么b的形状就是(1,3,3)

B.如果指定了num_classes，此时最后一维就等于num_classes，那么b的形状就是(1,3,4)

其实说白了最后一维就等于num_classes，区别只在于num_classes是否被提前指定而已