tensorflow.js

Printing the value of argMax index?


I would like to print the actual value of the argMax index (the highest probability) in this scenario:

const result = await model.predict(t4d); 
result.print(); // puts out: Tensor [[0.9899636, 0.0100364],]
result.as1D().argMax().print(); // prints either 0 or 1

Besides the index, I would want to print the actual value 0,XXXX behind argMax(). Any advice for this?

Test that doesn't work:

const confidence = result.dataSync<'float32'()>;
console.log(confidence);

Sorry if this is a question that has been answered any times, I have spent a few hours searching!


Solution

  • The value of the argMax index can be retrieved after getting the data of the tensors

    const result = await model.predict(t4d);
    const index = await result.as1D().argMax().data()[0]
    const predict = await result.data()
    
    // get the value
    const value = result[index]
    

    The other possibility would be to use topk

    const topk = result.as1D().topk()
    // get the highest value
    value = topk.value() // the value here is a tensor