Skip to content

[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 轉換成兩類的數值 —— 當然,多類其實也是可以的。

Leave a Reply取消回覆

Exit mobile version