axon

<h1><img src="https://yellow-cdn.veclightyear.com/35dd4d3f/60ecc516-97cf-4029-8792-61c74c6b5979.png" alt="Axon" width="350"></h1>

基于 Nx 的 Elixir 神经网络。

Axon 包括以下组件：

• Functional API – 数值定义 (defn) 的低级别 API，所有其他 API 都建立在此基础之上。
• Model Creation API – 用于管理模型初始化和应用的高级别模型创建 API。
• Training API – 受 PyTorch Ignite 启发的快速训练模型 API。

Axon 提供了抽象，使得在保持各组件之间独立的同时，能够轻松集成。你应该可以独立使用任何 API，而不依赖于其他 API。通过解耦 API，Axon 让你能够完全控制神经网络创建和训练的每个方面。Axon 使用 Polaris 作为其优化 API。

概述

欲了解详细概述，请参见：Axon: Deep Learning in Elixir

Functional API

在最低级别，Axon 包含多个模块，提供深度学习中常用方法的函数实现：

• Axon.Activations – 元素级激活函数。
• Axon.Initializers – 模型参数初始化函数。
• Axon.Layers – 常见的深度学习层实现。
• Axon.Losses – 常见的损失函数。
• Axon.Metrics – 例如准确率、绝对误差、精确度等训练指标。

函数 API 中的所有方法都以数值定义 (defn) 实现。这意味着你可以使用任何 Nx 编译器或后端加速 Axon。此外，你还可以将 Axon 函数 API 中的方法与自己的数值定义任意组合。Axon 完全在 Nx 张量上工作，因此构建在 Nx 之上的任何库都可能与 Axon 很好地集成。

因为 Axon 的高级别 API 构建在函数 API 之上，所以同样的优点也适用。每个神经网络都可以使用任何 Nx 编译器或后端即时编译 (JIT) 或提前编译 (AOT)，甚至转换为高层神经网络格式，如 TensorFlow Lite 和 ONNX。

模型创建

一个示例模型如下所示：

model =
  Axon.input("input", shape: {nil, 784})
  |> Axon.dense(128)
  |> Axon.dense(10, activation: :softmax)

模型只是一个 Elixir 结构体，因此将其序列化为多种格式在未来是很简单的。默认的 inspect 协议提供了模型的简单摘要。你可以使用 Axon.Display 模块来可视化一个更好地摘要。例如，你可以使用 Axon.Display.as_table/2 查看模型的表格摘要：

+-----------------------------------------------------------------------------------------------------------+
|                                                   Model                                                   |
+==================================+=============+==============+===================+=======================+
| Layer                            | Input Shape | Output Shape | Options           | Parameters            |
+==================================+=============+==============+===================+=======================+
| input ( input )                  | []          | {1, 784}     | shape: {nil, 784} |                       |
|                                  |             |              | optional: false   |                       |
+----------------------------------+-------------+--------------+-------------------+-----------------------+
| dense_0 ( dense["input"] )       | [{1, 784}]  | {1, 128}     |                   | kernel: f32[784][128] |
|                                  |             |              |                   | bias: f32[128]        |
+----------------------------------+-------------+--------------+-------------------+-----------------------+
| dense_1 ( dense["dense_0"] )     | [{1, 128}]  | {1, 10}      |                   | kernel: f32[128][10]  |
|                                  |             |              |                   | bias: f32[10]         |
+----------------------------------+-------------+--------------+-------------------+-----------------------+
| softmax_0 ( softmax["dense_1"] ) | [{1, 10}]   | {1, 10}      |                   |                       |
+----------------------------------+-------------+--------------+-------------------+-----------------------+
Total Parameters: 101770
Total Parameters Memory: 407080 bytes

Axon 提供了一些用于处理模型的方便工具。首先，我们选择了一种理念，即模型的唯一关心点是初始化和应用。这意味着模型不应该关心训练等细节。Axon 提供了 Axon.build/2 函数，用于将 Axon 数据结构构建为初始化和预测函数：

model =
  Axon.input("input", shape: {nil, 784})
  |> Axon.dense(128, activation: :relu)
  |> Axon.dropout(rate: 0.5)
  |> Axon.dense(10, activation: :softmax)

{init_fn, predict_fn} = Axon.build(model, compiler: EXLA)

params = init_fn.(Nx.template({1, 784}, :f32), %{})
predict_fn.(params, input)

你可以将函数直接传递给 defn，这意味着你可以轻松地将模型执行与现有数值定义集成。

目前，Axon 支持类似于 PyTorch 或 TensorFlow Keras 等框架中常见的高级别层。我们的目标是保持一个生产力高、可扩展、与其他现代深度学习框架齐平的 API。如果有你需要的功能没有包括，请随时提出问题。

优化和训练

训练 API 的目的是提供便捷和常规操作来实现训练循环。该 API 受优秀的 PyTorch Ignite 库启发。

训练模型的一般模式是：

1. 定义模型
2. 使用其中一个工厂方法 (此处为 Axon.Loop.trainer/3) 定义循环
3. 使用指标和事件处理程序调整循环
4. 在数据上运行循环

model =
  Axon.input("input", shape: {nil, 784})
  |> Axon.dense(128)
  |> Axon.dense(10, activation: :softmax)

model_state =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(0.005))
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.handle(:iteration_completed, &log_metrics/1, every: 50)
  |> Axon.Loop.run(data, %{}, epochs: 10, compiler: EXLA)

Axon 使用 Polaris 作为其优化 API。需要指出的是，优化 API 不直接依赖于 Axon 模型。你可以使用该 API 优化任何可微目标函数。

未来我们计划支持分布式训练循环。我们还在寻求通过在本地加速器上完全运行训练循环的方法来提高性能。

安装

要使用 Axon，你需要安装 Elixir。然后通过 mix 构建工具创建一个 Elixir 项目：

$ mix new my_app

然后将 Axon 添加到你的依赖项中：

def deps do
  [
    {:axon, "~> 0.6"}
  ]
end

你可能还需要包含一个 Nx 编译器，如 EXLA，以处理任何实际的深度学习工作负载：

def deps do
  [
    {:axon, "~> 0.6"},
    {:exla, "~> 0.6"},
  ]
end

与其他平台的集成

请参见 Ortex，它通过 ONNX Runtime 绑定提供完整的运行 ONNX 模型的兼容性。或者，参见 AxonONNX，将 ONNX 模型转换为 Axon 模型以在可能的情况下更好地与 Nx 集成。

赞助商

<a href="https://dockyard.com"><img src="https://yellow-cdn.veclightyear.com/35dd4d3f/71a9313b-7bcc-4a3a-956b-24c114d2dc6e.png" width=200 alt="DockYard"></a>

许可

版权所有 (c) 2021 Sean Moriarity

根据 Apache 许可证 2.0 版 (以下简称“许可证”) 授权，你不能使用此文件，除非遵守许可证。你可以获取许可证的副本，访问：http://www.apache.org/licenses/LICENSE-2.0

除非适用法律要求或书面同意，按“原样”分发的软件，许可证没有任何明示或暗示的担保条件。参见许可证了解管理权限和限制的详细内容。