Nx,Axonで始めるゼロから作るディープラーニング  7章畳み込みニューラルネットワーク

# はじめに
本記事はElixirで機械学習/ディープラーニングができるようになるnumpy likeなライブラリ [Nx](https://github.com/elixir-nx/nx/tree/main/nx)を使って
「[ゼロから作るDeep Learning ―Pythonで学ぶディープラーニングの理論と実装](https://www.oreilly.co.jp/books/9784873117584/)」
をElixirで書いていこうという記事になります。

今回は心が折れたのと、Nxで実装されたDeepLearningフレームワークの[Axon](https://github.com/elixir-nx/axon)が発表されたのでそれを使ってシンプルなCNNを実装していきます

書籍もtrainerとかブラックボックスなim2colとか出してきてるしもういいよね・・・？

[準備編](https://qiita.com/the_haigo/items/1a2f0b371a3644960251)
[exla setup](https://qiita.com/the_haigo/items/d1d1b62bfb9fd2013c17)
1章 pythonの基本 -> とばします
2章 パーセプトロン -> とばします
3章 [ニューラルネットワーク](https://qiita.com/the_haigo/items/bedd466142aaaf01641c)
[with exla](https://qiita.com/the_haigo/items/4ec791c878fd48b9dec1)
４章 [ニューラルネットワークの学習](https://qiita.com/the_haigo/items/2b7535bcff8b4df50d9d)
５章 [誤差逆伝播法](https://qiita.com/the_haigo/items/f273a9c687302507d79d)
[Nx.Defn.Kernel.grad](https://qiita.com/thehaigo/items/593926922a1d72591c3d)
6章 学習に関するテクニック -> とばします
7章 [畳み込みニューラルネットワーク](https://qiita.com/thehaigo/items/a4be7812b97bdddd6c1c/)

# Axon とは
Nxで書かれたNNライブラリでKerasのようなインターフェースを持っています

### model構築
denseのweightの初期値はglorot_uniformで生成されます
https://github.com/elixir-nx/axon/blob/e450f32416179baf818c2948e082932b475d8ed9/lib/axon.ex#L126
第2引数はニューロンの数(units)です

```elixir
model =
  Axon.input({nil, 784})
  |> Axon.dense(128)
  |> Axon.dense(10, activation: :softmax)
```

### 学習
trainerがあるので2行で簡単！
callbackやmetricsオプションはまだみたいですがroadmapに含まれているので気長に待ちましょう
それかTraining.trainのコードをコピーしてmetricsやcallbackを独自に組み込みましょう！

```elixir
model =
  Axon.input({nil, 784})
  |> Axon.dense(128)
  |> Axon.dense(10, activation: :softmax)

trained_params =
  model
  |> Axon.Training.step(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005))
  |> Axon.Training.train(x_test, t_test, epochs: 10, compiler: EXLA)
```

### 推論
学習したパラメーターとモデルを渡して推論します

```elixir
Axon.predict(model, trained_params, input, compiler: EXLA)
```

### import/export
roadmapに含まれてますが未実装ですので、学習済みパラメータをNx.to_flat_listでlistにしてcsvに書き込むなりetsにぶっこむなりするといいでしょう


### Layer
Axon.{Dence, Conv}等重み初期化等がついているのもありますが、単純に値の動きだけを見たい場合は
Axon.Layers以下にあるのでそちらを使いましょう
https://github.com/elixir-nx/axon/blob/main/lib/axon/layers.ex


# Axon のインストール
nx,exlaがbuildできているなら問題なくインストールできます

```deps.exs
defmodule NxDl.MixProject do
  ...
  def deps do
    [
      {:axon, "~> 0.1.0-dev", github: "elixir-nx/axon", branch: "main"}, #追加
      {:exla, "~> 0.1.0-dev", github: "elixir-nx/nx", sparse: "exla", override: true},
      {:nx, "~> 0.1.0-dev", github: "elixir-nx/nx", sparse: "nx", override: true},
      {:expyplot, "~> 1.1.2"},
      {:erlport, "~> 0.9.8" },
      {:benchee, "~> 1.0", only: :dev}
    ]
  end
end

```
まだ0.1.0-devで更新が頻繁に起こっているので、github経由のはdeps.update --allを実行しましょう

```
mix deps.get
mix deps.update --all
mix deps.compile
```

# 7.5 CNNの実装

### 学習データ
cnnのinputはデータ数,チャンネル数、幅、高さの4次元データである必要があるのでreshpeでチャンネル数の項目を追加します

```simple_conv_net.ex
  def train_inputs do
    x_train =
      Dataset.train_image
      |> Nx.tensor
      |> Nx.reshape({60000, 1, 28, 28})
      |> (& Nx.divide(&1, Nx.reduce_max(&1))).()
      |> Nx.to_batched_list(100)

    t_train =
      Dataset.train_label
      |> Dataset.to_one_hot
      |> Nx.tensor
      |> Nx.to_batched_list(100)
    {x_train, t_train}
  end

  def test_inputs do
    x_test =
      Dataset.test_image
      |> Nx.tensor
      |> Nx.reshape({10000,1,28,28})
      |> (& Nx.divide(&1, Nx.reduce_max(&1))).()

    t_test = Dataset.test_label |> Dataset.to_one_hot |> Nx.tensor
    {x_test, t_test}
  end
```

### model構築
畳み込み層１、Pooling層１のCNNなので以下のようになります
Axon.conv はunits数の箇所がfilterの数に当たります
他のパラメーターはデフォルトでstride 1, padding valid(0)となっています
書籍とは違ってPooling層のあとにFlatten層でcnnの形式のshapeをDenseで扱えるようにしています

```simple_conv_net.ex
  def model do
    Axon.input({nil,1,28,28})
    |> Axon.conv(30, kernel_size: {5, 5}, activation: :relu)
    |> Axon.max_pool(kernel_size: {2, 2})
    |> Axon.flatten()
    |> Axon.dense(100, activation: :relu)
    |> Axon.dense(10, activation: :softmax)
  end
```

### train & test
trainはstepでlossとoptimizerをセットして
trainでepoch数とcompilerをセットしています

testはまだtrainにmetricsの実行が組み込まれていないので、trained_paramsを引数にpredictを行ってその結果を正解ラベルと比較しています
またpredictを実行する場合はrequire Axonを書く必要があります
https://github.com/elixir-nx/axon/blob/main/examples/xor.exs

```simple_conv_net.ex
  require Axon

  def train do
    {x_train, t_train} = train_inputs()
    {trained_params, _optmizer} =
      model()
      |> Axon.Training.step(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005))
      |> Axon.Training.train(x_train, t_train, epochs: 10, compiler: EXLA)

    trained_params
  end

  def test(params) do
    {x_test, t_test} = test_inputs()
    Axon.predict(model(), params, x_test, compiler: EXLA)
    |> Axon.Metrics.accuracy(t_test)
  end
```

```bash
iex(1)> params = SimpleConvNet.train                                                        
Epoch 1, batch 600 of 600 - Average Loss: 4.9356760978698735

Epoch 1 Time: 9.746866s
Epoch 1 Loss: 4.935676097869873


Epoch 2, batch 600 of 600 - Average Loss: 4.6949691772460945

Epoch 2 Time: 9.08468s
Epoch 2 Loss: 4.694969177246094


Epoch 3, batch 600 of 600 - Average Loss: 4.6642417907714845

Epoch 3 Time: 9.071895s
Epoch 3 Loss: 4.664241790771484


Epoch 4, batch 600 of 600 - Average Loss: 4.6500945091247565

Epoch 4 Time: 9.131318s
Epoch 4 Loss: 4.650094509124756


Epoch 5, batch 600 of 600 - Average Loss: 4.6417756080627445

Epoch 5 Time: 9.053126s
Epoch 5 Loss: 4.641775608062744


Epoch 6, batch 600 of 600 - Average Loss: 4.6354269981384285

Epoch 6 Time: 9.057393s
Epoch 6 Loss: 4.635426998138428


Epoch 7, batch 600 of 600 - Average Loss: 4.6314291954040535

Epoch 7 Time: 9.047488s
Epoch 7 Loss: 4.631429195404053


Epoch 8, batch 600 of 600 - Average Loss: 4.6279087066650395

Epoch 8 Time: 9.081987s
Epoch 8 Loss: 4.627908706665039


Epoch 9, batch 600 of 600 - Average Loss: 4.6265153884887695

Epoch 9 Time: 9.098026s
Epoch 9 Loss: 4.6265153884887695


Epoch 10, batch 600 of 600 - Average Loss: 4.6234269142150885

Epoch 10 Time: 9.087951s
Epoch 10 Loss: 4.623426914215088

iex(2)> SimpleConvNet.test(params)  
#Nx.Tensor<
  f32
  0.9879000186920166
>

```
# 7.6 CNNの可視化

Nx.to_heatmapでみれますが、matrexと違って出力を貼っても表示れなくてお見せできないのが残念です

```
trained_params |> Tuple.to_list |> List.first |> Nx.to_heatmap
```
# EXLAをGPU modeで動かす
config.exsでexlaでcudaを使うように指定すると hostからcudaを使うようになります

```config.exs
use Mix.Config

config :exla, :clients, default: [platform: :cuda], cuda: [platform: :cuda]
```

# 最後に
実装は以上になります 
Axonはとてもスッキリしてて書いていて気持ちいいです
examplesもresnet50やfashionmnist autoencoder,mnist GANなど結構充実していて参考になります
RNN系やAttentionのlayerが実装されてTransformerとかが使えるようになるのが楽しみです!（使いこなせるかどうかは置いておいて）

最後まで読んでいただきありがとうございました

今回のコード
https://github.com/thehaigo/nx_dl/commit/65b920796c55fbaf7db397515ce3be3ff3b8b750

# 参考ページ
https://twitter.com/sean_moriarity/status/1380124787318665218
https://seanmoriarity.com/2021/04/08/axon-deep-learning-in-elixir/
https://github.com/elixir-nx/axon
https://github.com/elixir-nx/axon/blob/e450f32416179baf818c2948e082932b475d8ed9/lib/axon.ex#L126
https://github.com/elixir-nx/axon/blob/main/lib/axon/layers.ex
https://github.com/elixir-nx/axon/blob/main/examples/xor.exs
https://github.com/elixir-nx/axon/blob/main/examples/cifar10.exs