Search code examples
tensorflowtensorflow.js

Output probability of prediction in tensorflow.js


I have a model.json generated from tensorflow via tensorflow.js coverter

In the original implementation of model in tensorflow in python, it is built like this:

model = models.Sequential([
        base_model,
    layers.Dropout(0.2),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(num_classes)
    ])

In tensorflow, the probability can be generated by score = tf.nn.softmax(predictions[0]), according to the tutorial on official website.

How do I get this probability in tensorflow.js?

I have copied the codes template as below:

$("#predict-button").click(async function () {
    if (!modelLoaded) { alert("The model must be loaded first"); return; }
    if (!imageLoaded) { alert("Please select an image first"); return; }
    
    let image = $('#selected-image').get(0);
    
    // Pre-process the image
    console.log( "Loading image..." );
    let tensor = tf.browser.fromPixels(image, 3)
        .resizeNearestNeighbor([224, 224]) // change the image size
        .expandDims()
        .toFloat()
        // RGB -> BGR
    let predictions = await model.predict(tensor).data();
    console.log(predictions);
    let top5 = Array.from(predictions)
        .map(function (p, i) { // this is Array.map
            return {
                probability: p,
                className: TARGET_CLASSES[i] // we are selecting the value from the obj
            };
        }).sort(function (a, b) {
            return b.probability - a.probability;
        }).slice(0, 2);
        console.log(top5);
    $("#prediction-list").empty();
    top5.forEach(function (p) {
        $("#prediction-list").append(`<li>${p.className}: ${p.probability.toFixed(6)}</li>`);
        });

How should I modify the above code?

The output is just the same as the value of variable 'predictions':

Float32Array(5)
0: -2.5525975227355957
1: 7.398464679718018
2: -3.252196788787842
3: 4.710395812988281
4: -4.636396408081055
buffer: (...)
byteLength: (...)
byteOffset: (...)
length: (...)
Symbol(Symbol.toStringTag): (...)
__proto__: TypedArray


0: {probability: 7.398464679718018, className: "Sunflower"}
1: {probability: 4.710395812988281, className: "Rose"}
length: 2
__proto__: Array(0)

Please help!!! Thanks!


Solution

  • In order to extract the probabilities from the logits of the model using a softmax function you can do the following:

    This is the array of logits that are also the predictions you get from the model

    const logits = [-2.5525975227355957, 7.398464679718018, -3.252196788787842, 4.710395812988281, -4.636396408081055]
    

    You can call tf.softmax() on the array of values

    const probabilities = tf.softmax(logits)
    

    Result:

    [0.0000446, 0.9362511, 0.0000222, 0.0636765, 0.0000056]
    

    Then if you wanted to get the index with the highest probability you can make use of tf.argMax():

    const results = tf.argMax(probabilities).dataSync()[0]
    

    Result:

    1
    

    Edit

    I am not too familiar with jQuery so this might not be correct. But here is how I would get the probabilities of the outputs in descending order:

    let probabilities = tf.softmax(predictions).dataSync();
    $("#prediction-list").empty();
    probabilities.forEach(function(p, i) {
      $("#prediction-list").append(
        `<li>${TARGET_CLASSES[i]}: ${p.toFixed(6)}</li>`
      );
    });