Python PyTorch one_hot用法及代码示例

本文简要介绍python语言中 torch.nn.functional.one_hot 的用法。

用法:

torch.nn.functional.one_hot(tensor, num_classes=- 1) → LongTensor

参数：

• tensor(LongTensor) -任何形状的类值。

• num_classes(int) -类总数。如果设置为 -1，则类数将被推断为比输入张量中的最大类值大 1。

返回：

LongTensor 多了一个维度，在输入指示的最后一个维度的索引处值为 1，其他位置均为 0。

获取形状为 (*) 的索引值的 LongTensor 并返回形状为 (*, num_classes) 的张量，该张量到处都是零，除了最后一个维度的索引与输入张量的相应值匹配的情况下，在这种情况下它将为 1。

另请参阅One-hot on Wikipedia。

例子

>>> F.one_hot(torch.arange(0, 5) % 3)
tensor([[1, 0, 0],
        [0, 1, 0],
        [0, 0, 1],
        [1, 0, 0],
        [0, 1, 0]])
>>> F.one_hot(torch.arange(0, 5) % 3, num_classes=5)
tensor([[1, 0, 0, 0, 0],
        [0, 1, 0, 0, 0],
        [0, 0, 1, 0, 0],
        [1, 0, 0, 0, 0],
        [0, 1, 0, 0, 0]])
>>> F.one_hot(torch.arange(0, 6).view(3,2) % 3)
tensor([[[1, 0, 0],
         [0, 1, 0]],
        [[0, 0, 1],
         [1, 0, 0]],
        [[0, 1, 0],
         [0, 0, 1]]])