📑

Elixir/Nxを使ってFizzBuzzをディープラーニングで解く

2021/02/27に公開

Elixir

機械学習

fizzbuzz

tech

本記事では、ElixirとそのライブラリであるNxを使って、FizzBuzzをディープラーニングで解いてみたいと思います。ディープラーニングの定義についてはよく知らないので、突っ込まないでください。

ところてんさんによる以下の記事が元ネタです。楽しい問題のご提供、ありがとうございます。

Nxとは何か

Nxとは、Elixirで多次元のテンソルを効率的に計算でき、CPU/GPU/TPUによって実行できるコードを生成してくれるライブラリです。2021年2月18日に、Elixirの作者であるJosé Valimさんにより発表されました。

https://github.com/elixir-nx/nx

型付き・名前付きのテンソルを効率よく計算できる
リバースモードの自動微分を実装している
計算をCPU/GPU/TPUのバックエンドで実行できる
defnというマクロによりAOT/JITを用いたマルチステージのコンパイルができる

Pythonでいうと、NumPyとJAXに対応しており、Nxは両者をひとつのパッケージで扱えて便利とのことです。4については、あまりよくわかっていません。

使い方については、Joséさんによる以下のライブコーディング動画を観るのがよいと思います。

以下のスライドでは、その動画でライブコーディングされたコードを再現しつつ、簡単に試せるようにしてみたものについて説明しています。ご参照ください。

NxFizzBuzzの使い方

まずは、作ってみたコードについて見ていきたいと思います。NxFizzBuzzという名前をつけました。

https://github.com/kentaro/nx_fizzbuzz

使い方はREADMEにある通りです。

FizzBuzzのデータセットを生成
データセットを用いてモデルを学習させる
学習したパラメタを用いて、FizzBuzzの結果を推論する

この流れを、以下の通りのコードとして実行できます。

# generate a dataset
{features, labels} = NxFizzBuzz.Dataset.generate_dataset(10000, 20)

# train the model by the dataset
params = NxFizzBuzz.Model.fit(features, labels, epoch: 100, batch_size: 50, hidden_size: 8)

# predict answers
1..100
|> Enum.each(fn n ->
  NxFizzBuzz.predict_fizz_buzz(params, n)
  |> IO.puts()
end)

FizzBuzzとは何か

そもそも、FizzBuzzとは何でしょうか？Wikidediaでは、以下の通り説明されています。

プレイヤーは円状に座る。最初のプレイヤーは「1」と数字を発言する。次のプレイヤーは直前のプレイヤーの次の数字に1を足した数字を発言していく。ただし、3で割り切れる場合は「Fizz」（Bizz Buzzの場合は「Bizz」）、5で割り切れる場合は「Buzz」、両者で割り切れる場合（すなわち15で割り切れる場合）は「Fizz Buzz」（Bizz Buzzの場合は「Bizz Buzz」）を数の代わりに発言しなければならない。発言を間違えた者や、ためらった者は脱落となる。

Fizz Buzz - Wikipedia

知らないひとはあまりいないでしょうが、いったんこの定義は忘れてください。

データセットだけが与えられていて、この問題自体は知らないという設定でFizzBuzz問題にチャレンジするというのが面白ポイントです。

FizzBuzzの中身を探る

元ネタのところてんさんの記事「機械学習でFizzBuzzを実現する」では、データセットの特徴を探るために、まずはランダムフォレストを用いて分類してみています。しかし、データセットをそのまま用いただけでは分類がうまくできませんでした。

そこで、データセットから特徴量を引き出すことを試みています。

特徴量をそのまま入力することはだめだということが分かったので、特徴量を変換することを試みる。fizzbuzzなる問題は何やら倍数が関係しているらしいので、元の値を残しつつ剰余演算をした特徴量を追加するようにしてみる。feature engineeringってやつよ。

NxFizzBuzzの実装でみると、以下の部分です。

def to_feature(nums, size) do
  nums
  |> Enum.map(fn n ->
    1..size
    |> Enum.reduce([], fn m, acc ->
      [rem(n, m) | acc]
    end)
  end)
  |> Nx.tensor()
  |> Nx.divide(size)
end

このコードでは、独立変数の元データとして生成されたnumsとしてわたされる2^32までの範囲からランダムに選ばれた整数の列に対して、それぞれ1からsizeで渡された範囲の整数で割った剰余として特徴量を作っています。こういうのをやっていくのが「特徴量エンジニアリング」というのだなあと、勉強になります。また、問題の本質の周辺へいい感じに迫っているバランスが、とても面白いポイントですね。

モデルを作成する

上記で得られた特徴量に変換された元データとそれらに対応するラベルから、モデルを学習していきます。

入力は、上記で作成した次元数の特徴量です。それらを、パラメタで指定した次元数の中間層につなげます。その結果にソフトマックス関数を適用することで分類します。FizzBuzz問題の場合は、（1）引数に渡された整数、（2）Fizz、（3）Buzz、（4）FizzBuzzの4種類になるので、4クラス分類問題ということになります。

元データのベクトルと重みパラメタのドット積を取り、それにバイアス項を足してやったものに対して、活性化関数として用いるロジスティック関数を適用したものを次の中間層につなぎます。中間層でも同様に重みパラメタとバイアス項による計算を行ったあと、ソフトマックス関数で推論結果を出力します。

損失関数には交差エントロピー誤差を用います。それによって得られた損失から勾配を計算し、学習率を乗じた値を引くことでパラメタの学習を進めていきます。与えられたデータセットに対して繰り返しこのプロセスを実行していくことで、このモデルはFizzBuzz問題について適切な推論ができるようになっていくはずです。

defmodule NxFizzBuzz.Model do
  import Nx.Defn

  @default_epoch 10
  @default_batch_size 50
  @default_learning_rate 0.01
  @default_hidden_size 64

  def fit(x, y, opts \\ []) do
    opts =
      opts
      |> Keyword.put_new(:epoch, @default_epoch)
      |> Keyword.put_new(:batch_size, @default_batch_size)
      |> Keyword.put_new(:learning_rate, @default_learning_rate)
      |> Keyword.put_new(:hidden_size, @default_hidden_size)

    {_, input_size} = Nx.shape(x)

    init_params = {
      Nx.random_uniform({input_size, opts[:hidden_size]}, 0.0, 1.0, names: [:input, :hidden]),
      Nx.random_uniform({opts[:hidden_size]}, 0.0, 1.0, names: [:hidden]),
      Nx.random_uniform({opts[:hidden_size], 4}, 0.0, 1.0, names: [:hidden, :output]),
      Nx.random_uniform({4}, 0.0, 1.0, names: [:output])
    }

    zip =
      Enum.zip(
        Nx.to_batched_list(x, opts[:batch_size]),
        Nx.to_batched_list(y, opts[:batch_size])
      )
      |> Enum.with_index()

    for e <- 1..opts[:epoch],
        {{x_batch, y_batch}, b} <- zip,
        reduce: init_params do
      params ->
        IO.puts("epoch #{e}, batch #{b}")
        update(params, x_batch, y_batch, opts)
    end
  end

  defn predict({w1, b1, w2, b2}, batch) do
    batch
    |> Nx.dot(w1)
    |> Nx.add(b1)
    |> Nx.logistic()
    |> Nx.dot(w2)
    |> Nx.add(b2)
    |> softmax()
  end

  defnp update({w1, b1, w2, b2} = params, x, y, opts \\ []) do
    {grad_w1, grad_b1, grad_w2, grad_b2} = grad(params, loss(params, x, y))

    {
      w1 - grad_w1 * opts[:learning_rate],
      b1 - grad_b1 * opts[:learning_rate],
      w2 - grad_w2 * opts[:learning_rate],
      b2 - grad_b2 * opts[:learning_rate]
    }
  end

  defnp loss({w1, b1, w2, b2}, x, y) do
    preds = predict({w1, b1, w2, b2}, x)
    -Nx.sum(Nx.mean(Nx.log(preds) * y, axes: [:output]))
  end

  defnp softmax(t) do
    Nx.exp(t) / Nx.sum(Nx.exp(t), axes: [:output], keep_axes: true)
  end
end

実行結果

examples/ディレクトリにサンプルコードを置いてあるので、実行してみました。10,000個のデータセットから100回繰り返し学習をさせてみました。約500秒ほどかかって学習が終わり、1〜100までの数字について、FizzBuzzの結果を推論させてみたところ、100%正解することができました。

$ time mix run examples/fizz_buzz.exs

(snip)

97: prediction: 97, answer: 97, matched?: true
98: prediction: 98, answer: 98, matched?: true
99: prediction: Fizz, answer: Fizz, matched?: true
100: prediction: Buzz, answer: Buzz, matched?: true
================
Accuracy: 1.0

________________________________________________________
Executed in  487.17 secs   fish           external
   usr time  477.32 secs  126.00 micros  477.32 secs
   sys time   48.84 secs  688.00 micros   48.84 secs

今回はMacBook Pro（2018）のCPU（2.7 GHz クアッドコアIntel Core i7）を用いて計算してみましたが、EXLAを使ってGPUで計算させると、めちゃくちゃ速く終わりそうです。また、今回はやりませんでしたが、損失と正解率をグラフでプロットしてみると学習の様子が目に見えるようになり、もっと楽しめそうです。

おわりに

FizzBuzzという謎の問題について、与えられたデータセットからElixirのNxライブラリを用いて解くことができました。

そもそも、FizzBuzzといいう問題を知らない状況でデータセットのみが与えられたという設定で取り組んできたのでこんなまだるっこしいことをやってきたのでした。一方で、FizzBuzzというのは与えられた引数に対して計算で決定的に答えが求まる問題であるので、こんなおおげさなしかけを用いて解くようなことではありません。なんでもかんでも機械学習するのではなく、問題の特性をよく把握することが大事だということがわかりますね。