[译] 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. 他们相等。
2. 其中一个等于 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的核心功能。你学到了：

1. 如何使用 Nx.tensor 和 Nx.from_binary 来创建一个 tensor。2. 如何使用一元，二元和聚合操作来处理 tensor 3. 如何使用 defn 和 Nx 的 grad 来实现梯度下降。4. 如何使用 EXLA 编译器来加速数值计算。

尽管本文覆盖了开始使用 Nx 所需的基础知识，但还是有很多需要学习的。我希望本文可以驱使你继续学习关于 Nx 的项目，并且找到独特的使用场景。Nx仍在早期，有很多激动人心的东西在前方。

原文: https://dockyard.com/blog/202...