[譯] Nx 入門 -- Sean Moriarity

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
end

Nx.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))
end

MSE 測量目標值和預測的平均方差。越接近，則 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
  }
end

grad 輸入你想要獲取梯度的引數，以及一個引數化的函式，在這裡就是損失函式。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}
end

init_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...

[譯] Nx 入門 -- Sean Moriarity

適應 Tensor

Tensor 操作

基礎線性迴歸

總結

相關文章