torch.nn.functional.one_hot()

参数

• tensor(LongTensor)，任意shape的类别值
• num_classes(int)，所有类别的数量。默认为-1，如果设为-1，num_classes=tensor中最大的值+1

定义

假设tensor的shape为( * )，函数返回shape为( *, num_classes)的tensor，返回的tensor中除了最后一维索引为tensor对应值位置为1，其它都为0。

示例

target1 = torch.tensor([1, 2, 3])
print(target1.shape)  # torch.Size([3])
print(target1)  # tensor([1, 2, 3])
target1 = torch.nn.functional.one_hot(target1, num_classes=5)
print(target1.shape)  # torch.Size([3, 5])
print(target1)
# tensor([[0, 1, 0, 0, 0],
#         [0, 0, 1, 0, 0],
#         [0, 0, 0, 1, 0]])
target2 = torch.tensor([1, 2, 3])
target2 = torch.nn.functional.one_hot(target2)
print(target2.shape)  # torch.Size([3, 4])
print(target2)
# tensor([[0, 1, 0, 0],
#         [0, 0, 1, 0],
#         [0, 0, 0, 1]])

参考

torch.nn.functional.one_hot — PyTorch 1.12 documentation