[PyTorch] 將 Sigmoid 的輸出設定閥值（threshold）並轉成二元值

Last Updated on 2021-05-17 by Clay

在使用 PyTorch 當中的 Sigmoid 當我們的激活函數時，比如說接在模型的最後一層當二元分類的輸出，畢竟 Sigmoid 可以將數值壓在 [0-1] 之間，我們只要設定個閥值 (Threshold) ，比如說 0.5，就可以將數值分成倆類。

對於 Sigmoid 函數有興趣的，也許可以參考我曾經寫過的這片文章：《Sigmoid function》

那麼今天，我就來簡單紀錄如何設定閥值，然後將 Sigmoid 出來的結果分成兩類吧！

分類

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

import torch
import torch.nn as nn

input = torch.Tensor([[-1, -2, -3], [1, 2, 3]])
print(input)

Output:

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

然後我們將這個 Tensor 丟入 Sigmoid 函式中，會看到輸出值被壓縮在 [0 - 1] 之間。

sigmoid = nn.Sigmoid()

output = sigmoid(input)
print(output)

Output:

tensor([[0.2689, 0.1192, 0.0474],
        [0.7311, 0.8808, 0.9526]])

然後我們設定我們分界的 Threshold，假設 Threshold = 0.5，代表在數值在 0.5 以上的數值都被分到類別 1，在 0.5 以下的都被分到數值 0。

threshold = torch.tensor([0.5])

results = (output>threshold).float()*1
print(results)

Output:

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

我們可以這樣簡單地將 Sigmoid 的 Output 轉換成兩類的數值 —— 當然，多類其實也是可以的。