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]
]
今回は以上です。