JavaScript エンジニア向け: 知識ゼロから tensorflow.js で機械学習入門

この週末で機械学習を勉強した結果として、JavaScript エンジニア向けにまとめてみる。

自分が数式見て何もわからん…となったので、できるだけ動いてるコードで説明する。動いてるコードみてから数式見たら、多少気持ちがわかる感じになった。

最初に断っておくが、特にJSを使いたい理由がないなら python で keras 使ったほうがいいと思う。tensorflow.js が生きる部分もあるが、学習段階ではそこまで関係ないため。

追記: 最初 0 < a < 1.0 0 < b < 1.0三角関数 Math.sin をとっていて、これだと三角関数の一部の値しか使っておらず、線形に近似できそうな値を吐いていたので、次のように変更して、データも更新した。

// 修正前
const fn = (a, b) => {
  const n = Math.cos(a) * b + Math.sin(b) * a;
  return a > b ? n : -n;
};
// 修正版
const fn = (a, b) => {
  const n = Math.cos(a * Math.PI * 2) * b + Math.sin(b * Math.PI * 2) * a;
  return a > b ? n : -n;
};

これにより、精度はやや落ちた。

概要

やりたいこと: ある関数の、入力と出力を模倣したい。

ニューラルネットを使うとどうなるか。

  • モデルのネットワークは何かしらの関数を近似する
  • 入力と出力は何かしらの行列(tensor)で表現される

ニューラルネットワークでは、ある入力のときの出力を見て、モデルに期待する値との誤差を教えると、ネットワークがその値を生成するように変化する

真のモデル関数 f = (a, b) => a + b が あったとき、モデルにはこの中身を教えず、その振る舞いを覚えさせたい。

入力値 [1, 3] を与えた際、 model.predict([[1, 3]])[[0]] を生成したとする。このときの期待する値 4 との差 -4 が model に伝えられると、4 を生成しやすいように、内部のネットワークの重み(weight, bias)が変化する。

ニューラルネットの中身がどうなってるか、とりあえず一番最初はブラックボックスだと思ってよい。なんかいい感じにやってくれる。

実際どう定義されているか、誤差からどうその重みが変化するかは、バックプロパゲーションなどで検索。

バックプロパゲーション - Wikipedia

考え方

つまり、何かしらの状態を行列の表現に落として、その結果をどう評価するか既にわかっていて、また入力値が適切な因子として出力に関与することができ、また十分にモデルの内部ネットワークの複雑度があれば、大量のデータを与えると値が収束していく。

なので、何も仮説無くデータを突っ込んでも、適切なデータは出ないかもしれない。適切な因果を持つ、入力と出力のペアがあれば、データを与えれば与えるほど、精度が高くなる。

tensorflow.js で関数の振る舞いを予測する

ここから実践編。

適当に、こういう関数をでっち上げた。何も参考にしていない。完全にこの場の思いつきである。

const fn = (a, b) => {
  const n = Math.cos(a * Math.PI * 2) * b + Math.sin(b * Math.PI * 2) * a;
  return a > b ? n : -n;
};

この関数は、与えられた a, b に対して結構複雑な振る舞いをする。正直自分でも頭の中でグラフ掛けないし、予測がつかない。三角関数を使ってるのと、a > b の比較を行っているので、非線形(直線ではない)振る舞いをする。

このとき、(a, b) の出力は無限にこの関数で生成できるので、それを教師データとして使って、誤差を伝えるように訓練する。

以下、そのときの tensorflow.js を使ったコード。30000 回実行してみた。

import "@babel/polyfill";
import * as tf from "@tensorflow/tfjs";
import "@tensorflow/tfjs-node-gpu"; // Intel Cuda の対応OSのみ有効

const fn = (a, b) => {
  const n = Math.cos(a * Math.PI * 2) * b + Math.sin(b * Math.PI * 2) * a;
  return a > b ? n : -n;
};

function buildModel() {
  const model = tf.sequential();
  model.add(
    tf.layers.dense({ units: 20, activation: "relu", inputShape: [2] })
  );
  model.add(
    tf.layers.dense({
      units: 20,
      activation: "relu",
      inputShape: [1]
    })
  );
  model.add(
    tf.layers.dense({
      units: 20,
      activation: "relu"
    })
  );
  model.add(
    tf.layers.dense({
      units: 1,
      activation: "linear"
    })
  );

  model.compile({ optimizer: "sgd", loss: "meanSquaredError" });
  return model;
}


function train(model, count = 30000) {
  // create data
  const inputs: number[][] = new Array(count)
    .fill(0)
    .map(_ => [Math.random(), Math.random()]);
  const answers: number[][] = inputs.map(x => [fn(x[0], x[1])]);

  const xs = tf.tensor2d(inputs);
  const ys = tf.tensor2d(answers);

  return model.fit(xs, ys, {
    epochs: 100,
    callbacks: {
      onEpochEnd: async (epoch: any, log: any) => {
        console.log(`Epoch ${epoch}: loss = ${log.loss}`);
      }
    }
  });
}

async function run() {
  const model = buildModel();

  await train(model);

  console.log("--- use trained model");
  new Array(10).fill(0).forEach(() => {
    const input = [Math.random(), Math.random()];
    const pred: any = model.predict(tf.tensor2d([input]));
    const real = fn(input[0], input[1]);
    console.log("in", input, "pred", pred.dataSync()[0], "real", real);
  });
}

run();

なんか dense だの relu だの sgd だの meanSquaredError だの epoch だのいろんな概念が出てくるが、一旦はおまじないとして、気になったらググる感じで。実際には大したことない、身構えたのが馬鹿らしくなるような、シンプルな定義が多かった。

30000 回実行した後、本当に訓練できたか確認するのに、10 回実行してみている。

その結果

...

eta=0.0 ================================> loss=0.01
2367ms 79us/step - loss=0.01
Epoch 97: loss = 0.008017721727490426
Epoch 99 / 100
eta=0.0 ================================> loss=0.01
2408ms 80us/step - loss=0.01
Epoch 98: loss = 0.007897231815196575
Epoch 100 / 100
eta=0.0 ================================> loss=0.00
2309ms 77us/step - loss=0.01
Epoch 99: loss = 0.008164751819210747

--- use trained model
in [ 0.3962178765532509, 0.7682131016008036 ] pred 1.0202405452728271 real 1.0042189244656015
in [ 0.5636312166233701, 0.3811451901186962 ] pred 0.015699267387390137 real 0.031779120347658674
in [ 0.29081139359722297, 0.6065326337942547 ] pred 0.35762453079223633 real 0.3342764533539708
in [ 0.7485062557383053, 0.9206256550585616 ] pred 0.2782783508300781 real 0.3666547430935932
in [ 0.5138481068330991, 0.04557842574637938 ] pred 0.1076730489730835 real 0.09974545365571702
in [ 0.9162600474832179, 0.837541216933819 ] pred -0.043701171875 real -0.056856406373799406
in [ 0.7299011854033057, 0.8539316678532747 ] pred 0.5823078751564026 real 0.6872769887438305
in [ 0.02929035990521056, 0.4067624521196973 ] pred -0.445804238319397 real -0.4160877981686353
in [ 0.7649246937480239, 0.9202537029736069 ] pred 0.1884562373161316 real 0.28126628339301984
in [ 0.21538151181392773, 0.8708798400126363 ] pred -0.08348429203033447real -0.03174978980219534

思っていたより高い精度が出た。

正直、自分の思いつきの関数がここまで高い精度で模倣されたことにビビっている。とくに a > b の正負の出力に関してはほぼ確実に当ててきている。

最初は a + bMath.sin(a + b) でやってみて、すぐ収束したのでこれなら無理じゃないかって関数を突っ込んでみてこの精度なので、結構感動していた。訓練時間は手元で 3 分ぐらい。

ちなみに、 tensorflow.js はブラウザ環境では WebGPU を使って訓練するので、下手な CPU 環境より速い。

追記: 200000回

26分掛けて200000回やってみた。誤差 0.008 から 0.003程度まで精度が上がった

eta=0.0 ================================> loss=0.00
15717ms 79us/step - loss=0.00
Epoch 96: loss = 0.002879959831906017
Epoch 98 / 100
eta=0.0 ================================> loss=0.00
15722ms 79us/step - loss=0.00
Epoch 97: loss = 0.003037940525643062
Epoch 99 / 100
eta=0.0 ================================> loss=0.00
15697ms 78us/step - loss=0.00
Epoch 98: loss = 0.0030612829002237413
Epoch 100 / 100
eta=0.0 ================================> loss=0.00
15749ms 79us/step - loss=0.00
Epoch 99: loss = 0.0031942650708649306
--- use trained model
in [ 0.894890579256173, 0.791992136427107 ] pred -0.24931591749191284 real -0.23845978737014883
in [ 0.33861845138126845, 0.14223981949273057 ] pred 0.19588708877563477real 0.18875113058190596
in [ 0.533338516856545, 0.0756469011023253 ] pred 0.1673297882080078 real 0.17006682242955357
in [ 0.8359911024259461, 0.577326973633683 ] pred -0.09570872783660889 real -0.093409575348611
in [ 0.05713099891468687, 0.27499415441712993 ] pred -0.2965472340583801real -0.31389426496282413
in [ 0.2842413125538461, 0.19573535679085463 ] pred 0.2440524697303772 real 0.22609192489185875
in [ 0.8549135971844535, 0.7277203717553344 ] pred -0.4215628504753113 real -0.40083794824237673
in [ 0.8235456930613112, 0.2048808799839672 ] pred 0.8972318172454834 real 0.8820155692003951
in [ 0.8671403644770563, 0.47572531583815847 ] pred 0.4355853497982025 real 0.45111927209825686
in [ 0.853777499934028, 0.09838342580279558 ] pred 0.5768846273422241 real 0.5544972469695729
✨  Done in 1592.47s.

応用

tensorflow 自体は、入力と出力の最適化をしてくれるだけで、実際にどういうモデルを構築するか、どういう入力を与えるかは、結局人間が考える。

モデルをどう作るかは、最初は検討もつかないかもしれないが、自分もよくわからなくていろんな人に聞いたが、結局うまくいったモデルを参考にするしかないとのこと。大抵は論文になっていて、その解説があって、そのネットワークの構造をコードに落とすと再現できる(らしい。自分もまだそこまでたどり着いてない)。

流行りのものだと、画像を入力にする時にいい感じに次元を絞るのが CNN、時系列データから最終的な結果を全体に伝搬させたいのが DQN、2 つのモデルを競わせたり検証させたりするのが GAN、という理解をしている(GAN はまだふわっとしている。これから実装して勉強する)

勉強するなら

今回は js で書いたが、正直 python の方が keras での実装例が多くて勉強は捗ると思う。jupyter やら matplotlib だの色々便利ツールがある。

自分が tensorflow.js で遊んでるのは、ブラウザゲームのゲーム AI を作りたくて、js と python で同じ環境を二回実装するのが嫌、という理由なので、そういう理由がなければ python で keras 使っとくのがいいと思う。

python で実行した tensorflow のモデルを、tensorflow.js でも使えるような jsonコンパイルできるので、訓練環境はどっちでもいい。

というのが一週間ぐらい勉強した成果。