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