I might be getting this completely wrong, but please let me first try to explain what I need, and then what's wrong.

I have a classification task. The training data has 50 different labels. The customer wants to differentiate the low probability predictions, meaning that, I have to classify some test data as "Unclassified / Other" depending on the probability (certainty?) of the model.

When I test my code, the prediction result is a numpy array. One example is:

[[-1.7862008  -0.7037363   0.09885322  1.5318055   2.1137428  -0.2216074
   0.18905772 -0.32575375  1.0748093  -0.06001111  0.01083148  0.47495762
   0.27160102  0.13852511 -0.68440574  0.6773654  -2.2712054  -0.2864312
  -0.8428862  -2.1132915  -1.0157436  -1.0340284  -0.35126117 -1.0333195
   9.149789   -0.21288703  0.11455813 -0.32903734  0.10503325 -0.3004114
  -1.3854568  -0.01692022 -0.4388664  -0.42163098 -0.09182278 -0.28269592
  -0.33082992 -1.147654   -0.6703184   0.33038092 -0.50087476  1.1643585
   0.96983343  1.3400391   1.0692116  -0.7623776  -0.6083422  -0.91371405

I'm then using numpy.argmax() to identify the correct label.

My question is, is it possible to define a threshold (say, 0.6), and then compare the probability of the argmax() element so that I can classify the prediction as "other" if the probability is less than the threshold value?

Edit 1:

We are using 2 different models. One is Keras, and the other is BertTransformer. We have no problem in Keras since it gives the probabilities so I'm skipping Keras model.

The Bert model is pretrained. Here is how it is generated:

def model(self, data):
        number_of_categories = len(data['encoded_categories'].unique())
        model = BertForSequenceClassification.from_pretrained(

        # model.cuda()

        return model

The output given above is the result of model.predict() method. We compare both models, Bert is slightly ahead, therefore we know that the prediction works just fine. However, we are not sure what those numbers signify or represent.

Here is the Bert documentation.


1 Answer 1


Your call to model.predict() is returning the logits for softmax. This is useful for training purposes.

To get probabilties, you need to apply softmax on the logits.

import torch.nn.functional as F
logits = model.predict()
probabilities = F.softmax(logits, dim=-1)

Now you can apply your threshold same as for the Keras model.

  • $\begingroup$ Can you elaborate on how logits are used for training? All of the models I've used apply loss to the output of the softmax layer, I've never seen anything that uses the raw logits for calculating loss $\endgroup$
    – A Tyshka
    Jun 28, 2022 at 18:01
  • $\begingroup$ @ATyshka: Both Tensorflow and Pytorch provide this feature as standard, and I expect most other NN libraries do as well. As the original question was about Pytorch, here's a link for the Pytorch loss function for binary cross entropy using the raw logits: pytorch.org/docs/stable/generated/… - the OP already was using this or similar, so I don't see any need to put that in the answer. Hope it helps though - if you are running loss on the ouput layer, it is less efficient (and possibly unstable), so switch if you can $\endgroup$ Jun 28, 2022 at 20:25

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .