Nx,Exla GPUをdockerでビルドしbenchを回してみた

[準備編](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/)

ExlaのGPUビルドがめんどくさいので提供されているdockerファイルからビルドすると楽だったので手順をまとめてみました

以下のページを参考にdockerとcudaをセットアップ
https://www.digitalocean.com/community/tutorials/how-to-install-and-use-docker-on-ubuntu-20-04-ja
https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker
https://github.com/elixir-nx/nx/tree/main/exla#building-with-docker

```
docker build --rm -t exla:cuda10.1 . -f exla/docker/Dockerfile.cuda
docker run -it \
  -v $PWD:$PWD \
  -e TEST_TMPDIR=$PWD/tmp/bazel_cache \
  -e EXLA_CACHE=$PWD/tmp/exla_cache \
  -e EXLA_FLAGS=--config=cuda \
  -e EXLA_TARGET=cuda \
  -w $PWD \
  --gpus=all \
  --rm exla:cuda10.1 bash
```

runtimeがうまく読み込めないのでcudaで固定します

```nx/exla/config/runtime.exs
import Config

target = "cuda"

config :exla, :clients, default: [platform: String.to_atom(target)], cuda: [platform: :cuda]

config :logger, :console,
  format: "\n$time [$level] $metadata $levelpad$message\n",
  metadata: [:domain, :file, :line]

```

```bash
mix run bench/softmax.exs
```

```bash
Operating System: Linux
CPU Information: Intel(R) Core(TM) i7-9700K CPU @ 3.60GHz
Number of Available Cores: 8
Available memory: 31.30 GB
GeForce GTX 1060 6GB
Elixir 1.11.2
Erlang 23.2.3

Name              ips    average deviation     median     99th %
xla jit-gpu f32 keep    7734.20    0.129 ms   ±9.94%    0.127 ms    0.20 ms
xla jit-gpu f64 keep    2665.88    0.38 ms   ±6.68%    0.37 ms    0.51 ms
xla jit-gpu f32       505.45    1.98 ms   ±1.55%    1.97 ms    2.14 ms
xla jit-cpu f32       505.07    1.98 ms   ±1.66%    1.97 ms    2.14 ms
xla jit-gpu f64       257.88    3.88 ms   ±1.40%    3.86 ms    4.12 ms
xla jit-cpu f64       256.64    3.90 ms   ±1.75%    3.87 ms    4.17 ms
xla aot-cpu f32       209.75    4.77 ms   ±4.34%    4.73 ms    6.03 ms
xla aot-cpu f64       71.75    13.94 ms   ±1.92%    13.90 ms    15.79 ms
elixir f32          3.76   266.02 ms   ±1.76%   266.66 ms   274.19 ms
elixir f64          3.66   273.09 ms   ±0.40%   273.12 ms   274.78 ms

Comparison: 
xla jit-gpu f32 keep    7734.20
xla jit-gpu f64 keep    2665.88 - 2.90x slower +0.25 ms
xla jit-gpu f32       505.45 - 15.30x slower +1.85 ms
xla jit-cpu f32       505.07 - 15.31x slower +1.85 ms
xla jit-gpu f64       257.88 - 29.99x slower +3.75 ms
xla jit-cpu f64       256.64 - 30.14x slower +3.77 ms
xla aot-cpu f32       209.75 - 36.87x slower +4.64 ms
xla aot-cpu f64       71.75 - 107.79x slower +13.81 ms
elixir f32          3.76 - 2057.42x slower +265.89 ms
elixir f64          3.66 - 2112.11x slower +272.96 ms

Memory usage statistics:

Name          Memory usage
xla jit-gpu f32 keep     1.46 KB
xla jit-gpu f64 keep     1.46 KB - 1.00x memory usage +0 KB
xla jit-gpu f32       1.51 KB - 1.03x memory usage +0.0469 KB
xla jit-cpu f32       1.48 KB - 1.02x memory usage +0.0234 KB
xla jit-gpu f64       1.51 KB - 1.03x memory usage +0.0469 KB
xla jit-cpu f64       1.48 KB - 1.02x memory usage +0.0234 KB
xla aot-cpu f32       0.35 KB - 0.24x memory usage -1.10938 KB
xla aot-cpu f64       0.35 KB - 0.24x memory usage -1.10938 KB
elixir f32       273362.77 KB - 187114.62x memory usage +273361.30 KB
elixir f64       273368.20 KB - 187118.34x memory usage +273366.74 KB
```