[PyTorch] Set the threshold of Sigmoid output and convert it to binary value

When using sigmoid function in PyTorch as our activation function, for example it is connected to the last layer of the model as the output of binary classification. After all, sigmoid can compress the value between 0-1, we only need to set a threshold, for example 0.5 and you can divide the value into two categories.

If you want to get more information of sigmoid function, you can refer: [Machine Learning] Sigmoid function introduction

Today I will record how to set the threshold, and then divide the results from sigmoid function to two categories.

Classify

Suppose we have the following 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.]])

The we use sigmoid function to process tensor, and we can see the output values is between 0-1.

sigmoid = nn.Sigmoid()

output = sigmoid(input)
print(output)

Output:

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

Then we set the threshold of our demarcation. Assuming Threshold = 0.5, it means that all values above 0.5 are classified into category 1, and those below 0.5 are classified into value 0.

threshold = torch.tensor([0.5])

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

Output:

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