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

準備編
exla setup
1章 pythonの基本 -> とばします
2章 パーセプトロン -> とばします
3章 ニューラルネットワーク
with exla
４章 ニューラルネットワークの学習
５章 誤差逆伝播法
Nx.Defn.Kernel.grad
6章 学習に関するテクニック -> とばします
7章 畳み込みニューラルネットワーク

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]

mix run bench/softmax.exs

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