Google Colaboratory 上で Elixir Nx の各種バックエンドを動かす

## はじめに

以前の記事で、 Google Colaboratory 上で Livebook を起動し、 Nx を各種バックエンドで動かしました

https://qiita.com/RyoWakabayashi/items/45edc8ddbccba139a5ba

https://livebook.dev/

Elixir や Nx 、 evision のバージョンも上がっているため、今回は最新版で改めて実行します

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

https://github.com/cocoa-xu/evision

実装したノートブックはこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/colab/Livebook.ipynb

https://github.com/RyoWakabayashi/elixir-learning/blob/main/colab/bench.livemd

## Google Colaboratory 上に Elixir と Livebook をインストールする

Elixir と Livebook の最新版を　Google Colaboratory 上で動かす手順については以下を参照してください

https://qiita.com/RyoWakabayashi/items/51ff0f44065e958fc490

ただし、 evision を使うため、 Livebook 起動前に Jupyter 上で libopencv-dev をインストールしておく必要があります

```
!sudo apt-get install -y libopencv-dev
```

## セットアップ

セットアップセルで以下のコードを実行し、必要なモジュールをインストールします

`system_env` で各種環境変数を設定し、 CUDA を使うように指定しています

```elixir
Mix.install(
  [
    {:benchee, "~> 1.1"},
    {:nx, "~> 0.5"},
    {:exla, "~> 0.5"},
    {:torchx, "~> 0.5"},
    {:evision, "~> 0.1"},
    {:kino, "~> 0.8"}
  ],
  system_env: [
    {"XLA_TARGET", "cuda111"},
    {"LIBTORCH_TARGET", "cu113"},
    {"EVISION_ENABLE_CUDA", "true"},
    {"EVISION_CUDA_VERSION", "111"}
  ]
)
```

### EXLA で CUDA を使う

環境変数 `XLA_TARGET` に CUDA のバージョンを指定します

指定できる値は XLA の公式に書いてあります

https://github.com/elixir-nx/xla#xla_target

- cpu	
- tpu
- cuda118
- cuda114
- cuda111
- cuda
- rocm

`cuda118` が推奨ですが、 Google Colaboratory 上で指定すると何故か TorchX の方でエラーが出るようになったので、今回は `cuda111` を指定しています

### Torchx で CUDA を使う

環境変数 `LIBTORCH_TARGET` に CUDA のバージョンを指定します

指定できる値は以下のとおりです

- cpu
- cu102
- cu113
- cu116

:::note warn
2023/02/15 現在、 README の記載がおかしかったので Issue を出しました

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

### evision で CUDA を使う

かつては CUDA を使う場合コンパイルしないといけなかったのですが、 CUDA 用プリコンパイルを作ってくれたのでインストールにかかる時間が短縮されました

公式に従って、 `EVISION_ENABLE_CUDA=true` と `EVISION_CUDA_VERSION` に以下のいずれかを指定します

- 111
- 114
- 118

https://github.com/cocoa-xu/evision#evision_enable_cuda

## 各バックエンドでテンソルを生成する

### バイナリバックエンド

```elixir
tensor = Nx.tensor([1.0, 2.0, 3.0], backend: {Nx.BinaryBackend, device: :gpu})
Nx.add(tensor, tensor)
```

実行結果は以下のようになります

```elixir
#Nx.Tensor<
  f32[3]
  [2.0, 4.0, 6.0]
>
```

### EXLA バックエンド

```elixir
tensor = Nx.tensor([1.0, 2.0, 3.0], backend: {EXLA.Backend, device: :gpu})
Nx.add(tensor, tensor)
```

実行結果は以下のようになります

```elixir
#Nx.Tensor<
  f32[3]
  EXLA.Backend<cuda:0, 0.2185872277.2671902731.195814>
  [2.0, 4.0, 6.0]
>
```

### Torchx バックエンド

```elixir
tensor = Nx.tensor([1.0, 2.0, 3.0], backend: Torchx.Backend)
Nx.add(tensor, tensor)
```

実行結果は以下のようになります

```elixir
#Nx.Tensor<
  f32[3]
  Torchx.Backend(cuda)
  [2.0, 4.0, 6.0]
>
```

### evision バックエンド

```elixir
tensor = Nx.tensor([1.0, 2.0, 3.0], backend: Evision.Backend)
Nx.add(tensor, tensor)
```

実行結果は以下のようになります

```elixir
#Nx.Tensor<
  f32[3]
  Evision.Backend
  [2.0, 4.0, 6.0]
>
```

## ベンチマーク

200 x 200 の行列を足し算する処理でベンチマークします

```elixir
bench = fn backend ->
  tensor = Nx.iota({200, 200}, type: {:f, 64}, backend: backend)
  Nx.add(tensor, tensor)
end
```

```elixir
Benchee.run(%{
  "binary" => fn -> bench.({Nx.BinaryBackend, device: :gpu}) end,
  "exla" => fn -> bench.({EXLA.Backend, device: :gpu}) end,
  "torchx" => fn -> bench.(Torchx.Backend) end,
  "evision" => fn -> bench.(Evision.Backend) end
})
```

実行結果は以下のようになりました

```
Name              ips        average  deviation         median         99th %
torchx         8.26 K       0.121 ms    ±55.42%       0.110 ms        0.25 ms
exla           2.23 K        0.45 ms    ±37.34%        0.42 ms        0.79 ms
evision        0.87 K        1.14 ms    ±18.30%        1.09 ms        1.77 ms
binary       0.0557 K       17.96 ms    ±33.02%       15.02 ms       34.24 ms

Comparison: 
torchx         8.26 K
exla           2.23 K - 3.71x slower +0.33 ms
evision        0.87 K - 9.45x slower +1.02 ms
binary       0.0557 K - 148.23x slower +17.83 ms
```

|バックエンド|実行時間（最速との比）|
|:-------|---------------:|
|Torchx  |             最速|
|EXLA    |            3.71|
|Evision |            9.45|
|Binary  |          148.23|

## まとめ

Google Colaboratory 上で各種バックエンドを動かすことができました

ベンチマークはあくまでも単なる足し算の場合なので、参考までに