5_1_5 推論 その2 Tensor.argMax( axis? )メソッド

tf.Sequentialのpredict()が返すtf.Tensorオブジェクトは、推論結果の数値と結びついており、分類の場合、数値は確率と見なすことができます。

// 結果は確率の配列
predictOut.print(); // 例：[[0.1198563, 0.8801436],] -> versicolorの確率が高いと予測している

この結果の配列から、高い確率を取り出す方法として、Tensor.argMax( axis? )メソッド（tf.argMax (x, axis?)関数も同じ）の使用があります。argMax()は、TensorFlow.jsのドキュメンテーションによると、「Returns the indices of the maximum values along an axis(axis(軸)に沿った最大値のインデックスを返す)」とあります。これだけでは働きがよく分からないので、次のコードで見ていきます。

const a1 = tf.tensor2d([
    [5, 9, 2, 6],
    [4, 6, 7, 3],
    [9, 8, 3, 1]
], [3, 4]);

a1.print();
                    /*  [[5, 9, 2, 6],
                        [4, 6, 7, 3],
                        [9, 8, 3, 1]]
                    */
a1.argMax().print(); // [2, 0, 1, 0]
let axis = 0;
a1.argMax(axis).print(); // [2, 0, 1, 0]
axis = 1;
a1.argMax(axis).print(); // [1, 2, 0]

変数a1は、3行4列の配列を持つTensorオブジェクトです。a1はargMax()メソッドを持っており、axisという数値の引数が渡せます。デフォルト値は0です。

a1からargMax()を呼び出すと、[2, 0, 1, 0]という結果が返ります。またaxis=0のa1.argMax(axis)からも同じ結果が返ります。
一方、axisを1に指定したa1.argMax(axis)は[1, 2, 0]を返します。

axis値の違いによってargMax()が返す結果の違いの説明は、文字よりも図で示した方が分かりやすいです。

したがって、前の例のつづきで、上から３つめのボタンをクリックし[[0.1198563, 0.8801436],]というpredictOutが得られた場合には、横向きのaxis=1を使うことで、最大値のインデックス番号1が得られます。

 let axis = 0;
 predictOut.argMax(axis).print(); // [0, 0] -> axis=0は、縦向きに読む
 axis = 1;
 predictOut.argMax(axis).print(); // [1] -> axis=1は、横向きに読む

 // 予測結果の配列を横向きに読んだ場合での、最大値のインデックス位置を得る
 const predictOutMaxIndex = predictOut.argMax(axis);
 predictOutMaxIndex.print(); // [1]
 // その値
 const predictOutMaxIndexData = predictOutMaxIndex.dataSync()[0];
 console.log(predictOutMaxIndexData) // 1