LoginSignup
8
5
@sand(sand mypress)

Elixir Nx の基礎 (Livebook)

Last updated at Posted at 2022-11-08

Elixir Nx の基本的な概念や動作を、以下のドキュメントに沿って確認しました。Livebook 上で動作させています。
Introduction to Nx 公式ドキュメント

関連記事
Elixir Nx で学ぶ機械学習 - Axon無し (Livebook) - Qiita
Elixir Axon で学ぶ機械学習 (Livebook) - Qiita
Elixir Nx の基礎 (Livebook) - Qiita
Livebook 事始め (Elixir) - Qiita

1. Nx 概要

まず Livebook 上で必要なインストールを行っておきます。

Mix.install([
  {:nx, "~> 0.2"}
])

Nx は Elixir が効率よく数値計算を行うためのライブラリです。それは 型を持った多次元データ である テンソル 上で計算を行います。Nx は基本的に次の 3 つの能力を持っています。

  • Nx において テンソル は 型を持ったデータ であり、多次元であり、しかも次元には 名前 を与えることができる。
  • 数値関数の定義 (defn) で カスタムコードを関数化できます。それは tensor-aware オペレーションであり、関数です。
  • 自動微分 (Automatic differentiation) は、 autograd または autodiff として知られており、共通の数値計算のシナリオをサポートしています。機械学習 やシュミレーション, 曲線回帰、確率モデルなどです。

テンソルの定義は、型を持った多次元データである、ということです。単に、何重にも入れ子になった配列のイメージです。物理学で言う何らかの物理量を表すテンソルとは無関係です。

自動微分は機械学習やディープラーニングにおいて勾配計算を行う、最重要な概念です。

何故、微分 が必要になるかについては、興味があれば以下の過去記事を参照してください。
【機械学習】誤差逆伝播法のコンパクトな説明

Nx テンソルは以下の型を持ちます。

  • unsigned integers - u8, u16, u32, u64
  • signed integers - s8, s16, s32, s64
  • floats - f32, f64
  • brain floats - bf16
  • complex - c64, c128

2. Nx のテンソル

Nx のテンソルは、何ら物理量とは対応していないといいましたが、テンソルという言葉や概念は物理学由来だと思います。物理学ではスカラーは零階のテンソル、ベクトルは一階のテンソル、二階のテンソルと呼ぶらしいですが、ここでは次元という言葉を使っています。

  • 0次元テンソル - スカラー
  • 1次元テンソル - ベクトル
  • 2次元テンソル - 行列
  • 3次元テンソル - 行列の配列
  • ...
  • n次元テンソル - (n-1)次元テンソルの配列

2-1. テンソルの作成

Elixir リストからテンソルを作成することができます。以下のヘルプに詳細な例がありますが、いくつか見ていきましょう。

import IEx.Helpers
h Nx.tensor

0次元テンソル

    Nx.tensor(0)
    #Nx.Tensor<
      s64
      0
    >

1次元テンソル

    Nx.tensor([1, 2, 3])
    #Nx.Tensor<
      s64[3]
      [1, 2, 3]
    >

    Nx.tensor([1.2, 2.3, 3.4, 4.5])
    #Nx.Tensor<
      f32[4]
      [1.2000000476837158, 2.299999952316284, 3.4000000953674316, 4.5]
    >

型を明示的に指定できる。型指定で最大値を超えた場合はオーバフローとなる。
2進数(1,000,000 300) = 10進数(256)なので、 8ビットでオーバーフローを考慮すると、300 -256 = 44 になる。自動的に以下のような変換が行われる。

    Nx.tensor([300, 301, 302], type: :s8)
    #Nx.Tensor<
      s8[3]
      [44, 45, 46]
    >

2次元テンソル

    Nx.tensor([[1, 2, 3], [4, 5, 6]])
    #Nx.Tensor<
      s64[2][3]
      [
        [1, 2, 3],
        [4, 5, 6]
      ]
    >

これは以下のような行列に相当します。

\begin{pmatrix} 1 & 2 & 3\\ 4 & 5 & 6 \end{pmatrix}

3次元テンソル

Nx.tensor([[[1, 2], [3, 4], [5, 6]], [[-1, -2], [-3, -4], [-5, -6]]])
#Nx.Tensor<
  s64[2][3][2]
  [
    [
      [1, 2],
      [3, 4],
      [5, 6]
    ],
    [
      [-1, -2],
      [-3, -4],
      [-5, -6]
    ]
  ]
>

次元軸に名前を付ける

テンソル次元に名前を付けることができます。コードの可読性を高めてくれます。names は atom のリストですが、テンソルの多次元配列の入れ子の浅い軸の順に並びます。

    t1 = Nx.tensor([1, 2, 3], names: [:x])
    #Nx.Tensor<
      s64[x: 3]
      [1, 2, 3]
    >

    t2 = Nx.tensor([[1, 2, 3], [4, 5, 6]], names: [:x, :y])
    #Nx.Tensor<
      s64[x: 2][y: 3]
      [
        [1, 2, 3],
        [4, 5, 6]
      ]
    >

    t3 = Nx.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
      [[-1, -2, -3], [-4, -5, -6], [-7, -8, -9]]], names: [:batch, :height, :width])
    #Nx.Tensor<
      s64[batch: 2][height: 3][width: 3]
      [
        [
          [1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]
        ],
        [
          [-1, -2, -3],
          [-4, -5, -6],
          [-7, -8, -9]
        ]
      ]
    >

名前でテンソルの要素にアクセスしてみます。あわせて Nx.shape/1 でテンソルの shape を取得し、Nx.reshape/2 で shape を変更します。

t1[x: 1]
#Nx.Tensor<
  s64
  2
>

t2[x: 1]
#Nx.Tensor<
  s64[y: 3]
  [4, 5, 6]
>

t2[y: 1]
#Nx.Tensor<
  s64[x: 2]
  [2, 5]
>

t2[x: 0, y: 1]
#Nx.Tensor<
  s64
  2
>

t3[height: 1]
#Nx.Tensor<
  s64[batch: 2][width: 3]
  [
    [4, 5, 6],
    [-4, -5, -6]
  ]
>

t3[width: 2]
#Nx.Tensor<
  s64[batch: 2][height: 3]
  [
    [3, 6, 9],
    [-3, -6, -9]
  ]
>


Nx.shape(t1)
{3}

Nx.shape(t2)
{2, 3}

Nx.shape(t3)
{2, 3, 3}

Nx.reshape(t3, {6, 3}, names: [:batches, :values])
#Nx.Tensor<
  s64[batches: 6][values: 3]
  [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
    [-1, -2, -3],
    [-4, -5, -6],
    [-7, -8, -9]
  ]
>

バイナリ変換

Nx.tensor([[1, 2], [3, 4]], type: :u8) |> Nx.to_binary()
<<1, 2, 3, 4>>

Nx.from_binary(<<0, 1, 2>>, :u8)
#Nx.Tensor<
  u8[3]
  [0, 1, 2]
>

2-2. テンソル関数

テンソル関数とは任意の shape のテンソルを引数にとれる関数です。

Nx.sum(t3[0])
#Nx.Tensor<
  s64
  45
>

Nx.sum(t2)
#Nx.Tensor<
  s64
  21
>

Nx.sum(t2, axes: [:x])
#Nx.Tensor<
  s64[y: 3]
  [5, 7, 9]
>

tensor1 = Nx.tensor([[1, 2], [3, 4]])
tensor2 = Nx.tensor([[5, 6], [7, 8]])
Nx.subtract(tensor2, tensor1)
#Nx.Tensor<
  s64[2][2]
  [
    [4, 4],
    [4, 4]
  ]
>

2-3. ブロードキャスト

テンソルでの計算ですが、例えば以下のような引き算は、

\begin{pmatrix} 1 & 2 & 3\\ 4 & 5 & 6 \end{pmatrix} - 1 = \begin{pmatrix} 0 & 1 & 2\\ 3 & 4 & 5 \end{pmatrix}

数学的な意味では以下の引き算に置き換えられると、便利です。

\begin{pmatrix} 1 & 2 & 3\\ 4 & 5 & 6 \end{pmatrix} - \begin{pmatrix} 1 & 1 & 1\\ 1 & 1 & 1 \end{pmatrix} = \begin{pmatrix} 0 & 1 & 2\\ 3 & 4 & 5 \end{pmatrix}

Nx.broadcast/2 を使えば、1 を {2, 2} テンソルに変換してくれます。

Nx.broadcast(1, {2, 2})
#Nx.Tensor<
  s64[2][2]
  [
    [1, 1],
    [1, 1]
  ]
>

Nx.broadcast(1, tensor) は 1 を tensor の shape にブロードキャストしてくれます。また例えば Nx.subtract(tensor, 1) も2つの引数のテンソルの shape を互換性のあるものへと自動的にブロードキャストしてくれます。

[[1], [2]] |> Nx.tensor() |> Nx.broadcast({1, 2, 4})
#Nx.Tensor<
  s64[1][2][4]
  [
    [
      [1, 1, 1, 1],
      [2, 2, 2, 2]
    ]
  ]
>

以下のヘルプで多くの example が表示できます。

h Nx.broadcast

3. 数値関数の定義 (defn)

defn を使って好きなテンソル関数を作り出すことができます。引数として、スカラーやベクトル、行列などのテンソルを受けることができます。

引き算と足し算の関数を定義します。

defmodule TensorMath do
  import Nx.Defn

  defn hiku(a, b) do
    a - b
  end

  defn tasu(a, b) do
    a + b
  end
end

行列テンソル同士の引き算です。

tensor1 = Nx.tensor([[1, 1], [1, 1]])
tensor2 = Nx.tensor([[5, 6], [7, 8]])
TensorMath.hiku(tensor2, tensor1)
#Nx.Tensor<
  s64[2][2]
  [
    [4, 5],
    [6, 7]
  ]
>

tensor1 をスカラー1に変えても同じ結果です。

TensorMath.hiku(tensor2, 1)
#Nx.Tensor<
  s64[2][2]
  [
    [4, 5],
    [6, 7]
  ]
>

足し算についても同じ結果が得られます。

TensorMath.tasu(tensor2, tensor1)
#Nx.Tensor<
  s64[2][2]
  [
    [6, 7],
    [8, 9]
  ]
>

TensorMath.tasu(tensor2, 1)
#Nx.Tensor<
  s64[2][2]
  [
    [6, 7],
    [8, 9]
  ]
>

4. 自動微分 (autograd)

h Nx.Defn.grad

Nx.Defn.grad(fun) は無名関数を引数として、新たな無名関数を作成し返してくれます。その無名関数は与えられた地点の勾配を求めてくれます。

sinの微分はcosですので、0 地点の勾配は 1 になります。

    fun = Nx.Defn.grad(fn x -> Nx.sin(x) end)
    fun.(Nx.tensor(0))
    #Nx.Tensor<
      f32
      1.0
    >

以下のような f(x) 関数の導関数は g(x) となります。

f(x) = 3x^{2} + 2x + 1\\
g(x) = 6x + 2

この事実をもとに Nx の自動微分を行ってみます。
地点 2 の勾配を求めてみます。

defmodule Funs do
  import Nx.Defn

  defn poly(x) do
    3 * Nx.power(x, 2) + 2 * x + 1
  end

  defn poly_slope_at(x) do
    grad(&poly/1).(x)
  end
end

Funs.poly_slope_at(2)
#Nx.Tensor<
  f32
  14.0
>

g(2) = 14 の結果とあっています。

5. テンソルの軸 (Axis)

テンソルの軸は少しわかりにくい概念なので、もう少し述べておきたいと思います。
便利な関数なのですが、Nx.iota() に shape を与えると、0　始まりの連番の値を持つテンソルを作ってくれます。

Nx.iota({5})
#Nx.Tensor<
  s64[5]
  [0, 1, 2, 3, 4]
>

Nx.iota({3, 2, 3}, names: [:batch, :height, :width])
#Nx.Tensor<
  s64[batch: 3][height: 2][width: 3]
  [
    [
      [0, 1, 2],
      [3, 4, 5]
    ],
    [
      [6, 7, 8],
      [9, 10, 11]
    ],
    [
      [12, 13, 14],
      [15, 16, 17]
    ]
  ]
>

テンソルの軸はテンソルの多次元配列の深さを示すものですが、axis_index(tensor, axis) はテンソルの軸を与えることによって、テンソルの index を返してくれるものです。
-1 は一番最後の軸に相当します。

Nx.axis_index(Nx.iota({100, 10, 20}), 0)
0

Nx.axis_index(Nx.iota({100, 10, 20}), -1)
2

Nx.axis_index(Nx.iota({100, 10, 20}, names: [:batch, :x, :y]), :x)
1

axis_size(tensor, axis) は与えられた軸のサイズを返します。

Nx.axis_size(Nx.iota({100, 10, 20}), 0)
100

Nx.axis_size(Nx.iota({100, 10, 20}, names: [:batch, :x, :y]), :y)
20

new_axis(tensor, axis, name \ nil) はサイズ1 の軸を追加します。

t = Nx.tensor([[1, 2, 3], [4, 5, 6]])
#Nx.Tensor<
  s64[2][3]
  [
    [1, 2, 3],
    [4, 5, 6]
  ]
>

Nx.new_axis(t, 0, :new)
#Nx.Tensor<
  s64[new: 1][2][3]
  [
    [
      [1, 2, 3],
      [4, 5, 6]
    ]
  ]
>

Nx.new_axis(t, 1, :new)
#Nx.Tensor<
  s64[2][new: 1][3]
  [
    [
      [1, 2, 3]
    ],
    [
      [4, 5, 6]
    ]
  ]
>

Nx.new_axis(t, 2, :new)
#Nx.Tensor<
  s64[2][3][new: 1]
  [
    [
      [1],
      [2],
      [3]
    ],
    [
      [4],
      [5],
      [6]
    ]
  ]
>

Nx.new_axis(t, -1, :new)
#Nx.Tensor<
  s64[2][3][new: 1]
  [
    [
      [1],
      [2],
      [3]
    ],
    [
      [4],
      [5],
      [6]
    ]
  ]
>

（例）one-hot エンコーディング

one-hot エンコーディング とは、例えば 3 というスカラー値を、[0, 0, 0, 1, 0, 0, 0, 0, 0, 0] という 3番目だけ 1 で他は全て 0 のベクトルに変換することです。0 番始まりの 3 番目です。

わかり易いように段階を踏みます。

left =
  1..4
  |> Enum.to_list
  |> Nx.tensor
  |> Nx.new_axis(-1)
#Nx.Tensor<
  s64[4][1]
  [
    [1],
    [2],
    [3],
    [4]
  ]
>

right = Nx.tensor(Enum.to_list(0..9))
#Nx.Tensor<
  s64[10]
  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
>

Nx.equal/2 は2つの tensor を比べて、等しい箇所を１に、そうでない個所を０にします。

Nx.equal(left, right)
#Nx.Tensor<
  u8[4][10]
  [
    [0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 1, 0, 0, 0, 0, 0]
  ]
>

（補足）内部では以下のようなブロードキャストが行われた後で、Nx.equal で要素ごとの比較が行われているものと推測されます。

left = 
  [
    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
    [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
    [3, 3, 3, 3, 3, 3, 3, 3, 3. 3],
    [4, 4, 4, 4, 4, 4, 4, 4, 4, 4]
  ]

right =
  [
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],  
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],  
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],  
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
  ]

今回は以上です。

8
5
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
8
5