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.DefnendNx.Defn 模块中蕴含了 defn 的定义。它是一个可用于定义数值计算的宏。数值计算和 Elixir 函数的应用办法雷同，但仅反对一个无限的语言子集，为了反对 JIT。defn 还替换了很多Elixir 的外围办法，例如：
defn add_two(a, b) do a + bend+ 主动转换成了 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 enddata = 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 endend接着，咱们须要定义损失 （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 ) endend在训练循环里，咱们从 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 endendtarget_m = :rand.normal(0.0, 10.0)target_b = :rand.normal(0.0, 5.0)target_fn = fn x -> target_m * x + target_b enddata = 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.exsTarget m: -0.057762353079829236 Target b: 0.681480460783122Learned 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.exsTarget m: -1.4185910271067492 Target b: -2.9781437461823965Learned m: -1.4185925722122192 Learned b: -2.978132724761963Training time: 4.460695s咱们胜利实现了学习。这一次，花了 4.5 秒。当初，为了利用EXLA的JIT编译，将上面这个模块属性增加到你的模块中：
defmodule LinReg do import Nx.Defn @default_defn_compiler EXLAend它会通知Nx应用EXLA编译器来编译所有数值计算。当初，从新运行一遍：
Target m: -3.1572039775886167 Target b: -1.9610560589959405Learned m: -3.1572046279907227 Learned b: -1.961051106452942Training 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...