Nx 是一個 BEAM 上的，用於操作張量（tensor）和數值計算的新庫。Nx 期望為elixir、erlang以及其它 BEAM 語言開啟一扇大門，通往一個嶄新的領域 -- 使用者能夠使用 JIT 和高度特殊化的 tensor 操作來加速他們的程式碼。本文中，你會學到基礎的操作 Nx 的方法，以及如何將其用於機器學習應用中。
適應 Tensor
Nx 的 Tensor 類似於 PyTorch 或 TensorFlow 的 tensor，NumPy 的多維陣列。用過它們，那就好辦。不過它與數學定義不完全一致。Nx 從 Python 生態裡借鑑了許多，所以適應起來應該是很容易。Elixir 程式設計師可以把 tensor 想象為巢狀列表，附帶了一些後設資料。
iex> Nx.tensor([[1, 2, 3], [4, 5, 6]])
#Nx.Tensor<
s64[2][3]
[
[1, 2, 3],
[4, 5, 6]
]
>Nx.tensor/2 是用來建立 tensor 的，它可以接受巢狀列表和標量：
iex> Nx.tensor(1.0)
#Nx.Tensor<
f32
1.0
>後設資料在 tensor 被檢視時可以看到，比如例子裡的 s64[2][3] 和 f32。Tensor 有形狀和型別。每個維度的長度所組成的元祖構成了形狀。在上面的例子裡第一個 tensor 的形狀是 {2, 3}，表示為 [2][3]：
iex> Nx.shape(Nx.tensor([[1, 2, 3], [4, 5, 6]]))
{2, 3}把 tensor 想象為巢狀列表的話，就是兩個列表，每個包含3個元素。巢狀更多：
iex> Nx.shape(Nx.tensor([[[[1, 2, 3], [4, 5, 6]]]]))
{1, 1, 2, 3}1個列表，其包含1個列表，其包含2個列表，其包含3個元素。
這種思維在處理標量時可能會有點困惑。標量的形狀是空元組：
iex> Nx.shape(Nx.tensor(1.0))
{}因為標量是 0 維的 tensor。它們沒有任何維度，所以是“空”形。
Tensor 的型別就是其中數值的型別。Nx裡型別表示為一個二元元組，包含類與長度或位元寬度：
iex> Nx.type(Nx.tensor([[1, 2, 3], [4, 5, 6]]))
{:s, 64}
iex> Nx.type(Nx.tensor(1.0))
{:f, 32}型別很重要，它告訴 Nx 在內部應該如何儲存 tensor。Nx 的 tensor 在底層表示為binary：
iex> Nx.to_binary(Nx.tensor(1))
<<1, 0, 0, 0, 0, 0, 0, 0>>
iex> Nx.to_binary(Nx.tensor(1.0))
<<0, 0, 128, 63>>關於大端小端，Nx 使用的是硬體本地的端序。如果你需要 Nx 使用指定的大小端，你可以提一個 issue 來描述使用場景。
Nx 會自動判斷輸入的型別，你也可以指定某種型別：
iex> Nx.to_binary(Nx.tensor(1, type: {:f, 32}))
<<0, 0, 128, 63>>
iex> Nx.to_binary(Nx.tensor(1.0))
<<0, 0, 128, 63>>因為 Nx tensor 內部表示是 binary，所以你不應該使用 Nx.tensor/2，它在創造特別大的 tensor 時會非常昂貴。Nx 提供了 Nx.from_binary/2 這個方法，不需要遍歷巢狀列表：
iex> Nx.from_binary(<<0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64>>, {:f, 32})
#Nx.Tensor<
f32[3]
[1.0, 2.0, 3.0]
>Nx.from_binary/2 輸入一個 binary 和型別，返回一個一維的 tensor。如果你想改變形狀，可以用 Nx.reshape/2：
iex> Nx.reshape(Nx.from_binary(<<0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64>>, {:f, 32}), {3, 1})
#Nx.Tensor<
f32[3][1]
[
[1.0],
[2.0],
[3.0]
]
>reshape 只是改變了形狀屬性，所以是想當便宜的操作。當你有binary格式的資料，使用 from_binary 在 reshape 是最高效的做法。
Tensor 操作
如果你是 Elixir 程式設計師，一定很熟悉 Enum 模組。因此，你可能會想要使用 map 和 reduce 方法。Nx 提供了這些方法，但你應當不去使用它們。
Nx 裡的所有操作都是 tensor 相關的，即它們可用於任意形狀和型別的 tensor。例如，在 Elixir 裡你可能習慣這樣做：
iex> Enum.map([1, 2, 3], fn x -> :math.cos(x) end)
[0.5403023058681398, -0.4161468365471424, -0.9899924966004454]但在 Nx 裡你可以這樣：
iex> Nx.cos(Nx.tensor([1, 2, 3]))
#Nx.Tensor<
f32[3]
[0.5403022766113281, -0.416146844625473, -0.9899924993515015]
>Nx 裡所有的一元操作都是這樣 -- 將一個函式應用於 tensor 裡的所有元素：
iex> Nx.exp(Nx.tensor([[[1], [2], [3]]]))
#Nx.Tensor<
f32[1][3][1]
[
[
[2.7182817459106445],
[7.389056205749512],
[20.08553695678711]
]
]
>
iex> Nx.sin(Nx.tensor([[1, 2, 3]]))
#Nx.Tensor<
f32[1][3]
[
[0.8414709568023682, 0.9092974066734314, 0.14112000167369843]
]
>
iex> Nx.acosh(Nx.tensor([1, 2, 3]))
#Nx.Tensor<
f32[3]
[0.0, 1.316957950592041, 1.7627471685409546]
>幾乎沒必要使用 Nx.map，因為對元素的一元操作總是可以達到相同的效果。Nx.map 總是會低效一些，而且你沒法使用類似 grad 的變換。此外，一些 Nx 後端和編譯器不支援 Nx.map，所以可移植性也是問題。Nx.reduce 也是一樣。使用 Nx 提供的聚合方法，類似 Nx.sum, Nx.mean, Nx.product 是比 Nx.reduce 更好的選擇：
iex> Nx.sum(Nx.tensor([1, 2, 3]))
#Nx.Tensor<
s64
6
>
iex> Nx.product(Nx.tensor([1, 2, 3]))
#Nx.Tensor<
s64
6
>
iex> Nx.mean(Nx.tensor([1, 2, 3]))
#Nx.tensor<
f32
2.0
>Nx 聚合方法還支援在單個軸上的聚合。例如，如果你有一攬子樣本，你可能只想計算單個樣本的均值：
iex> Nx.mean(Nx.tensor([[1, 2, 3], [4, 5, 6]]), axes: [1])
#Nx.Tensor<
f32[2]
[2.0, 5.0]
>甚至給定多個軸：
iex> Nx.mean(Nx.tensor([[[1, 2, 3], [4, 5, 6]]]), axes: [0, 1])
#Nx.Tensor<
f32[3]
[2.5, 3.5, 4.5]
>Nx 還提供了二元操作。例如加減乘除：
iex> Nx.add(Nx.tensor([1, 2, 3]), Nx.tensor([4, 5, 6]))
#Nx.Tensor<
s64[3]
[5, 7, 9]
>
iex> Nx.subtract(Nx.tensor([[1, 2, 3]]), Nx.tensor([[4, 5, 6]]))
#Nx.Tensor<
s63[1][3]
[-3, -3, -3]
>
iex> Nx.multiply(Nx.tensor([[1], [2], [3]]), Nx.tensor([[4], [5], [6]]))
#Nx.Tensor<
s64[3][1]
[
[4],
[10],
[18]
]
>
iex> Nx.divide(Nx.tensor([1, 2, 3]), Nx.tensor([4, 5, 6]))
#Nx.Tensor<
f32[3]
[0.25, 0.4000000059604645, 0.5]
>二元操作有一個限定條件，那就是tensor 的形狀必須能廣播到一致。在輸入的 tensor 形狀不同時會觸發廣播：
iex> Nx.add(Nx.tensor(1), Nx.tensor([1, 2, 3]))
#Nx.Tensor<
s64[3]
[2, 3, 4]
>這裡，標量被廣播成了更大的 tensor。廣播可以讓我們實現更節約記憶體的操作。比如，你想把一個 50x50x50 的tensor 乘以 2，你可以直接藉助廣播，而不需要創造另一個全是 2 的 50x50x50 的 tensor。
廣播的兩個 tensor 的每個維度必須是匹配的。當符合下列條件中的一個時，維度就是匹配的：
- 他們相等。
- 其中一個等於 1.
當你試圖廣播不匹配的 tensor 時，會遇到如下報錯：
iex> Nx.add(Nx.tensor([[1, 2, 3], [4, 5, 6]]), Nx.tensor([[1, 2], [3, 4]]))
** (ArgumentError) cannot broadcast tensor of dimensions {2, 3} to {2, 2}
(nx 0.1.0-dev) lib/nx/shape.ex:241: Nx.Shape.binary_broadcast/4
(nx 0.1.0-dev) lib/nx.ex:2430: Nx.element_wise_bin_op/4如果需要的話，你可以用 expanding, padding, slicing 輸入的 tensor 來解決廣播問題；但要小心。
基礎線性迴歸
目前為止，我們只是在 iex 裡面學習簡單的例子。我們所有的例子都可以被 Enum 和列表來實現。本小節，我們要展現 Nx 真正的力量，使用梯度下降來解決基礎線性迴歸問題。
建立一個新的 Mix 專案，包含 Nx 和它的後端。在這裡，我會使用 EXLA, 你也可以使用 Torchx。他們有一些區別，但都可以執行下面的例子。
def deps do
[
{:exla, "~> 0.2"},
{:nx, "~> 0.2"}
]
end然後執行:
$ mix deps.get && mix deps.compile第一次執行可能需要一段時間的下載和編譯，你可以在 EXLA 的 README 裡找到一些提示。
當 Nx 和 EXLA 都編譯好後，建立一個新檔案 regression.exs。在其中建立一個模組：
defmodule LinReg do
import Nx.Defn
endNx.Defn 模組中包含了 defn 的定義。它是一個可用於定義數值計算的宏。數值計算和 Elixir 函式的使用方法相同，但僅支援一個有限的語言子集，為了支援 JIT。defn 還替換了很多Elixir 的核心方法，例如：
defn add_two(a, b) do
a + b
end+ 自動轉換成了 Nx.add/2。defn 還支援特殊變換：grad。grad 宏會返回一個函式的梯度。梯度反映了一個函式的變化率。細節這裡就不提了，現在，只需要掌握如何使用 grad。
如上所述，我們將使用梯度下降來解決基本線性迴歸問題。線性迴歸是對輸入值和輸出值之間的關係進行建模。輸入值又稱為解釋值，因為它們具有解釋輸出值的因果關係。舉個實際的例子，你想透過日期、時間、是否有彈窗來預測網站的訪問量。你可以收集幾個月以來的資料，然後建立一個基礎迴歸模型來預測日均訪問量。
在我們的例子中，我們將會建立一個有一個輸入值的模型。首先，在LinReg 模組之外定義我們的訓練集:
target_m = :rand.normal(0.0, 10.0)
target_b = :rand.normal(0.0, 5.0)
target_fn = fn x -> target_m * x + target_b end
data =
Stream.repeatedly(fn -> for _ <- 1..32, do: :rand.uniform() * 10 end)
|> Stream.map(fn x -> Enum.zip(x, Enum.map(x, target_fn)) end)
IO.puts("Target m: #{target_m}\tTarget b: #{target_b}\n")首先，我們定義了 target_m,target_b,target_fn。我們的線性方程是 y = m*x +b，所以我們使用 Stream 來重複生成了一攬子輸入輸出對。我們的目標是使用梯度下降來學習 target_m 和 target_b。
接下來我們要定義的是模型。模型是一個引數化的函式，將輸入轉化為輸出。我們知道我們的函式格式是 y = m*x + b，所以可以這樣定義：
defmodule LinReg do
import Nx.Defn
defn predict({m, b}, x) do
m * x + b
end
end接著，我們需要定義損失 （loss） 函式。Loss 函式通常用來測量預測值和真實值的誤差。它能告訴你模型的優劣。我們的目標是最小化loss函式。
對於線性迴歸問題，最常用的損失函式是均方誤差 mean—squared error (MSE)：
defn loss(params, x, y) do
y_pred = predict(params, x)
Nx.mean(Nx.power(y - y_pred, 2))
endMSE 測量目標值和預測的平均方差。越接近，則 MSE 越趨近於零。我們還需要一個方法來更新模型，使得 loss 減小。我們可以使用梯度下降。它計算 loss 函式的梯度。梯度能告訴我們如何更新模型引數。
一開始很難講清楚梯度下降在做什麼。想象你正在尋找一個湖的最深處。你有一個測量儀在船上，但沒有其它資訊。你可以搜查整個湖，但這會耗費無限的時間。你可以每次在一個小範圍裡找到最深的點。比如，你測量出往左走深度從5變成7，往右走深度從5變成3，那麼你應該往左走。這就是梯度下降所做的，給你一些如何改變引數空間的資訊。
你可以透過計算損失函式的梯度，來更新引數：
defn update({m, b} = params, inp, tar) do
{grad_m, grad_b} = grad(params, &loss(&1, inp, tar))
{
m - grad_m * 0.01,
b - grad_b * 0.01
}
endgrad 輸入你想要獲取梯度的引數，以及一個引數化的函式，在這裡就是損失函式。grad_m 和 grad_b 分別是 m 和 b 的梯度。透過將 grad_m 縮小到 0.01 倍，再用 m 減去這個值，來更新 m。這裡的 0.01 也叫學習指數。我們想每次移動一小步。
update返回更新後的引數。在這裡我們需要m和b的初始值。在尋找深度的例子裡，想象你有一個朋友知道最深處的大概位置。他告訴你從哪裡開始，這樣我們能夠更快地找到目標：
defn init_random_params do
m = Nx.random_normal({}, 0.0, 0.1)
b = Nx.random_normal({}, 0.0, 0.1)
{m, b}
endinit_random_params 隨機生成均值 0.0 方差 0.1 的引數 m 和 b。現在你需要寫一個訓練迴圈。訓練迴圈輸入幾捆樣本，並且應用 update，直到某些條件達到時才停止。在這裡，我們將10次訓練 200 捆樣本：
def train(epochs, data) do
init_params = init_random_params()
for _ <- 1..epochs, reduce: init_params do
acc ->
data
|> Enum.take(200)
|> Enum.reduce(
acc,
fn batch, cur_params ->
{inp, tar} = Enum.unzip(batch)
x = Nx.tensor(inp)
y = Nx.tensor(tar)
update(cur_params, x, y)
end
)
end
end在訓練迴圈裡，我們從 stream 中提取200捆資料，在每捆資料後更新模型引數。我們重複epochs次，在每次更新後返回引數。現在，我們只需要呼叫 LinReg.train/2 來返回學習後的 m 和 b:
{m, b} = LinReg.train(100, data)
IO.puts("Learned m: #{Nx.to_scalar(m)}\tLearned b: #{Nx.to_scalar(b)}")總之，regression.exs 現在應該是：
defmodule LinReg do
import Nx.Defn
defn predict({m, b}, x) do
m * x + b
end
defn loss(params, x, y) do
y_pred = predict(params, x)
Nx.mean(Nx.power(y - y_pred, 2))
end
defn update({m, b} = params, inp, tar) do
{grad_m, grad_b} = grad(params, &loss(&1, inp, tar))
{
m - grad_m * 0.01,
b - grad_b * 0.01
}
end
defn init_random_params do
m = Nx.random_normal({}, 0.0, 0.1)
b = Nx.random_normal({}, 0.0, 0.1)
{m, b}
end
def train(epochs, data) do
init_params = init_random_params()
for _ <- 1..epochs, reduce: init_params do
acc ->
data
|> Enum.take(200)
|> Enum.reduce(
acc,
fn batch, cur_params ->
{inp, tar} = Enum.unzip(batch)
x = Nx.tensor(inp)
y = Nx.tensor(tar)
update(cur_params, x, y)
end
)
end
end
end
target_m = :rand.normal(0.0, 10.0)
target_b = :rand.normal(0.0, 5.0)
target_fn = fn x -> target_m * x + target_b end
data =
Stream.repeatedly(fn -> for _ <- 1..32, do: :rand.uniform() * 10 end)
|> Stream.map(fn x -> Enum.zip(x, Enum.map(x, target_fn)) end)
IO.puts("Target m: #{target_m}\tTarget b: #{target_b}\n")
{m, b} = LinReg.train(100, data)
IO.puts("Learned m: #{Nx.to_scalar(m)}\tLearned b: #{Nx.to_scalar(b)}")現在你可以這樣執行：
$ mix run regression.exs
Target m: -0.057762353079829236 Target b: 0.681480460783122
Learned m: -0.05776193365454674 Learned b: 0.6814777255058289看我們的預測結果是多麼地接近！我們已經成功地使用梯度下降來實現線性迴歸；然而我們還可以更進一步。
你應該注意到了，100個epochs的訓練花費了一些時間。因為我們沒有利用EXLA提供的JIT編譯。因為這是個簡單的例子，然而，當你的模型變得複雜，你就需要JIT的加速。首先，我們來看一下EXLA和純elixir在時間上的區別：
{time, {m, b}} = :timer.tc(LinReg, :train, [100, data])
IO.puts("Learned m: #{Nx.to_scalar(m)}\tLearned b: #{Nx.to_scalar(b)}\n")
IO.puts("Training time: #{time / 1_000_000}s")在沒有任何加速的情況下：
$ mix run regression.exs
Target m: -1.4185910271067492 Target b: -2.9781437461823965
Learned m: -1.4185925722122192 Learned b: -2.978132724761963
Training time: 4.460695s我們成功完成了學習。這一次，花了 4.5 秒。現在，為了利用EXLA的JIT編譯，將下面這個模組屬性新增到你的模組中：
defmodule LinReg do
import Nx.Defn
@default_defn_compiler EXLA
end它會告訴Nx使用EXLA編譯器來編譯所有數值計算。現在，重新執行一遍：
Target m: -3.1572039775886167 Target b: -1.9610560589959405
Learned m: -3.1572046279907227 Learned b: -1.961051106452942
Training time: 2.564152s執行的結果相同，但時間從4.5s縮短到了2.6s，幾乎60%的提速。必須承認，這只是一個很簡單的例子，而你在複雜的實現中看到的速度提升遠不止這些。比如，你可以試著實現MNIST，一個epoch使用純elixir將花費幾個小時，而EXLA會在0.5s～4s左右完成，取決於你的機器使用的加速器。
總結
本文覆蓋了Nx的核心功能。你學到了：
- 如何使用 Nx.tensor 和 Nx.from_binary 來建立一個 tensor。2. 如何使用一元，二元和聚合操作來處理 tensor 3. 如何使用 defn 和 Nx 的 grad 來實現梯度下降。4. 如何使用 EXLA 編譯器來加速數值計算。
儘管本文覆蓋了開始使用 Nx 所需的基礎知識，但還是有很多需要學習的。我希望本文可以驅使你繼續學習關於 Nx 的專案，並且找到獨特的使用場景。Nx仍在早期，有很多激動人心的東西在前方。
原文: https://dockyard.com/blog/202...