torch.nn.functional.one_hot

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

接受形状为 (*) 的包含索引值的 LongTensor，并返回形状为 (*, num_classes) 的张量，该张量在所有位置都为零，除了最后一维的索引与输入张量的对应值匹配的位置，该位置的值为 1。

另请参见 Wikipedia 上的独热编码 。

参数

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

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

返回值

LongTensor，在最后一维的索引由输入指示的位置具有 1 值，其他位置为 0，该张量在最后一维具有一个额外的维度。

示例

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