[PyTorch] 將 Tensor 型態轉換成 One-Hot Encoding 型態

Last Updated on 2021-05-27 by Clay

在處理進行 Machine Learning 的資料時，我有著『將 Labels 轉成 One-Hot Encoding 型態』這樣的需求。我本來想得很單純，就將 Tensor 轉成 Numpy，再由 Numpy 轉成 One-Hot —— 就像我在這篇《在 Numpy 中將數值轉成 One-Hot 型態》中講述的一樣。

但後來我發現不對；與其轉成 Numpy、轉成 One-Hot、再轉回 Tensor，不如打從一開始就直接將 Tensor 轉成 One-Hot Encoding 的型態，省去型態轉換造成的浪費。

有相關需求的人也可以直接參考這篇： https://discuss.pytorch.org/t/convert-int-into-one-hot-format/507

以下，我就來紀錄該怎麼在 Tensor 當中轉換成 One-Hot 吧！

One-Hot

在開始前，我簡單講解一下 One-Hot 的型態，這裡不討論為何要轉換，只單純敘述 One-Hot 是怎麼樣的編碼。

假設我們有以下這一組數值：

[1, 2, 3]

那我們將其轉換成 One-Hot 編碼後：

[[0, 1, 0, 0],
 [0, 0, 1, 0],
 [0, 0, 0, 1]]

Index 從 0 開始計算的話， One-Hot Encoding 就是上述的型態。

使用 scatter_() 來進行轉換

假設我們有以下這樣一組的 Tensor：

y = torch.tensor([[1], [2], [3]])
print(y)

Output:

tensor([[1],
        [2],
        [3]])

如果是這樣的 Label 型態，代表著我們在 Training 的過程中每次是以 batch_size = 3 的尺寸丟 Training data 訓練 —— 這感覺相當少見，不過這裡就讓我們這樣假設吧！

然後我們初始化一個符合這個 batch_size 的全零 Tensor：

batch_size = 3
length = 4
y_onehot = torch.zeros([batch_size, length])
print(y_onehot)

Output:

tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])

由於 Index 都是從 0 開始計算，故我們有 index = 3 的 Label 時，我們會有 length = 4 的矩陣大小。

然後來到最重要的一步了：我們使用 scatter_() 來轉換 One-Hot:

print(y_onehot.scatter_(1, y, 1))

Output:

tensor([[0., 1., 0., 0.],
        [0., 0., 1., 0.],
        [0., 0., 0., 1.]])

這樣我們就有了 One-Hot Encoding 的 Tensor 了！