Nx を Evision.Backend で動かす

## はじめに

Evision に Nx バックエンドが実装されました

試してみたいと思います

:::note info
2023/08/02

Evision の 0.1.33 がリリースされたため、更新
:::

## 実装コード

バックエンドの実装はこちら

https://github.com/cocoa-xu/evision/blob/main/lib/evision/backend.ex

## CPU で実行

ノートブックの全量はこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/evision/backend.livemd

### CPU 実行環境

- MacBook Pro 13 inch, 2019
  - CPU 2.4 GHz クアッドコアIntel Core i5
  - メモリ 16 GB
- Elixir 1.15.4
- Erlang 26.0.2
- Livebook 0.10.0

### セットアップ

まずセットアップします

他のバックエンドと比較するため、 EXLA や Torchx、 Benchee を入れています

```elixir
Mix.install([
  {:benchee, "~> 1.1"},
  {:nx, "~> 0.5"},
  {:exla, "~> 0.5"},
  {:torchx, "~> 0.5"},
  {:evision, "~> 0.1.33"},
  {:kino, "~> 0.10"}
])
```

### バックエンドの動作確認

各バックエンドが動作することを確認します

バイナリバックエンド

```elixir
tensor = Nx.tensor([1, 2, 3], type: :f64, backend: Nx.BinaryBackend)
Nx.add(tensor, tensor)
```

![スクリーンショット 2022-10-26 19.25.32.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/6b6fa7ea-cd72-c4af-ee2f-df5abd8731b4.png)

EXLA バックエンド

```elixir
tensor = Nx.tensor([1, 2, 3], type: :f64, backend: EXLA.Backend)
Nx.add(tensor, tensor)
```

![スクリーンショット 2022-10-26 19.26.20.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/b46014f8-66d4-9899-0c0e-25f849236b35.png)

Torchx バックエンド

```elixir
tensor = Nx.tensor([1, 2, 3], type: :f64, backend: Torchx.Backend)
Nx.add(tensor, tensor)
```

![スクリーンショット 2022-10-26 19.27.06.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/61c736f0-6ab7-0782-b2aa-7441028e50f9.png)

Evision バックエンド

```elixir
tensor = Nx.tensor([1, 2, 3], type: :f64, backend: Evision.Backend)
Nx.add(tensor, tensor)
```

![スクリーンショット 2022-10-26 19.27.48.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/0e8deb7c-fd6f-d0cb-d308-0864d054d62c.png)

ちなみに、以下の場合はエラーになります

```elixir
tensor = Nx.tensor([1, 2, 3], backend: Evision.Backend)
Nx.add(tensor, tensor)
```

![スクリーンショット 2022-10-26 19.28.43.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/a92bf269-6ca2-5596-80ca-318705789699.png)

Evision.Backend が s64 型をサポートしていないためです

https://github.com/cocoa-xu/evision#unsupported-type-map

また、以下の場合もエラーになります

```elixir
tensor = Nx.tensor([1, 2, 3], type: :f64, backend: Evision.Backend)
Nx.concatenate([tensor, tensor])
```

![スクリーンショット 2022-10-26 19.30.39.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/4584bcdf-3463-732c-f0d6-05a172deb735.png)

これは Evision.Backend が実装していない関数だからです

### CPU 環境でのベンチマーク

比較するためにベンチマーク用の関数を用意します

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

実行してみます

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

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

![スクリーンショット 2023-08-02 9.24.29.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/205f7d7a-5ebb-a47d-9a28-1002aaf4b1c2.png)

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

## GPU で実行

以下の記事でやったように Google Colaboratory 上で Livebook を起動し、 GPU で動かしてみました

https://qiita.com/RyoWakabayashi/items/113b94866780c7646af1

### GPU環境のセットアップ

```elixir
Mix.install(
  [
    {:benchee, "~> 1.1"},
    {:nx, "~> 0.5"},
    {:exla, "~> 0.5"},
    {:torchx, "~> 0.5"},
    {:evision, "~> 0.1.33"},
    {:kino, "~> 0.10"}
  ],
  system_env: [
    {"XLA_TARGET", "cuda118"},
    {"EXLA_TARGET", "cuda"},
    {"LIBTORCH_TARGET", "cu116"},
    {"EVISION_ENABLE_CUDA", "true"},
    {"EVISION_ENABLE_CONTRIB", "true"},
    {"EVISION_CUDA_VERSION", "118"}
  ]
)
```

### GPU環境用ベンチマーク関数

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

### GPU環境ベンチマーク実行

Torchx は CPU も使えるため、 CPU 版も計測しています

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

### GPU環境ベンチマーク実行結果

![スクリーンショット 2023-08-02 12.43.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/1d32b65e-d3db-161c-81a6-2d90ec225f33.png)

|   バックエンド  | 実行時間（最速との比） |
|:------------|-----------------:|
| Torchx      |              最速 |
| EXLA        |             3.78 |
| Torchx(CPU) |             4.76 |
| Evision     |            13.35 |
| Binary      |           141.25 |

Torchx の CPU と GPU で結構差が出ていますね

## まとめ

CPU でも GPU でも Evision.Backend を動かすことができました

型や関数のサポートが一部ないことは気をつけないといけませんね