轴突：Elixir中的深度学习

脉冲的长螺纹部分，脉冲从细胞体传导到其他细胞。

今天，我很高兴能够公开宣布Axon，这是一个在Elixir中创建神经网络的图书馆。轴突仍然是预发布;但是，我相信它已达到了一个点，即可从开源社区进行实验和输入。在这篇文章中，我将介绍Axon的API，讨论一些设计决策，并制定了图书馆的未来计划。

Axon的主要目标之一是在易用性和灵活性之间取得平衡。一端，我们希望Depare Learning在深度学习体验中的程序员发现Axon易于使用和初学者。另一方面，我们希望经验丰富的从业者和研究人员发现Axon是一种富有成效的，灵活和令人耳目一新的功能，以深入的学习框架。

功能API - 所有其他API的Elixir Defn的低级API。

模型创建API - 管理模型初始化和应用程序的高级模型创建API。

从头开始，我们尝试创建抽象，以便在保持每个组件之间的分离级别时轻松集成。您应该能够在没有依赖性的情况下使用任何API。通过以这种方式解耦API，Axon可以完全控制创建和培训神经网络的每个方面。

在最低级别的情况下，Axon由许多模块组成，具有深度学习中的常用方法的功能实现：

此API在很大程度上与Tenorflow中的Pytorch或TF.NN功能很大程度上类似于Torch.nn.Functional。功能实现非常裸露，不具有在更高级别的API中提供的相同便利。缺乏便利性具有最大的灵活性和控制模型和模型培训的益处。

功能API中的所有方法都被实现为NX DEFN。这意味着您可以使用任何NX编译器或后端来加速Axon。此外，您可以在Axon功能API中任意撰写方法，使用您自己的数值定义。 Axon完全在NX Tensors上工作，所以任何内置在NX上的图书馆都可能与Axon相结合。

由于Axon的高级API在功能API之上构建，因此适用相同的优点。您可以使用任何NX编译器或后端来加速模型培训或推理，可能是AOT编辑的可能性。

模型创建API的目标是消除与创建，初始化和应用模型相关的大多数样板。此外，我们还希望以易于导出到其他格式（如Onnx或Tensorflow Lite）的方式构建和代表模型。 Axon使用Axon Struct解决这些问题：

Axon结构代表模型的计算图。 Axon模型创建API中的每种方法都接受Axon结构并返回新的Axon结构，其中输入图层表示模型创建中的“原子”操作。一个示例模型看起来像：

通过这种方法，您可以使用常规的Elixir函数来表示模型构建块，并无论如何您都认为适合：

因为底层模型只是一个Elixir struct，所以模型序列化比通过您想要的其他格式将Axon节点转换为等同的模型序列化并不难。目前，我们仍然习惯模型序列化的细节，并喜欢对潜在需求和用例的反馈。

使用常规ELIXIR STRUCT的增加的好处使用ICSPACT协议轻松定制模型检查：

Axon为使用型号提供了一些便利设施。首先，我们选择采取哲学，模型唯一担忧是初始化和应用。这意味着模型不应该关注所有培训细节。 Axon提供宏：axon.init / 2和axon.predict / 4用于初始化和应用模型：

Axon.init / 2和Axon.predict / 4也可以从任何地方使用 - 内部排污并嵌套，常规Elixir代码：

您应该注意，为了加速这些函数，您需要使用nx.defn.jit / 3。

模型API还可以轻松地使用NX层应用现有数值定义和NX代码：

有计划支持经常性层，注意层，还有更多。我们的目标是维持一个富有成效，可扩展的API和与其他现代深度学习框架的API。如果有功能，您需要查看路线图上不包含在路线图中，请随时打开一个问题。

Axon的型号优化API采用与DeepMind的Optax相同的方法。 API的目标是提供用于创建高级优化器的低级构造，然后提供基于该API之上的高级优化器。

Axon将优化器视为元组：{init_fn / 1，update_fn / 2}。 init_fn / 1接受模型的参数并初始化优化器状态。 Update_fn / 2接受“更新”（基于梯度的优化中最常见的渐变）和优化器状态，​​并返回变换更新和新的优化器状态。

在最低级别，Axon在Axon.updates中实现了许多更新函数。每个更新函数充当组合器 - 接受init_fn / 1和update_fn / 2的元组，因为它是第一个参数并返回修改的init_fn和update_fn的元组。这意味着您可以任意编写更新以构建复杂的优化器。作为示例，您可以使用Axon.updates API来实现adam优化程序，如：

如果您发现Axon.updates中的转换太高，则可以使用axon.updates.stateful / 3和axon.updates来实现自定义转换。 axon.updates.stateful / 3代表有状态转型：

值得注意的是，优化API不直接取决于Axon模型。您可以使用API​​优化任何可差异的目标函数。

在未来，我们计划支持与学习率计划的集成，并探索更高级的优化方法，包括差异私有的SGD和二阶方法。

培训API的目的是提供用于实施训练环的便利性和常用例程。在Dougal Maclaurin的（Autograd的创造者毗邻Matt Johnson和David Duvenaud）博士论文，他写道：

自身射出的目标是让渐变努力地。如果您可以编写丢失功能，Autograd应该能够为您提供梯度。

他还描述了在优化目标函数参数时获得梯度的有用性。基本上，Autograd简化了将机器学习算法编写到编写可分化目标函数的任务的任务。这是我们在Axon的训练API中维护的哲学。如果您可以编写参数化，可差异的目标函数并将其与数据配对，则可以使用Axon的培训API。

API部分受到优秀的Pytorch Lightning库的启发。在这篇文章时，Axon训练API由2种方法组成：Axon.Training.Step和Axon.tring.train。在实践中，您可以使用这些方法来培训这样的轴突模型：

axon.training.step代表了一个训练步骤。它返回一个元组：{init_fn，step_fn}，它们分别表示训练初始化函数和训练步骤函数。 Axon.tration.step / 3的形式在此示例中实际上只是Axon.training.Treat.step / 2接受两个元组的方便：

Objective_FN是一个参数化的可差异目标函数，可接受模型参数，输入和标签，并返回丢失。回想一下，Axon优化器也只是元组，因此init_update_fn和update_fn都具有与Axon Optimizer相同的形式。组合所有这些方法以产生在训练期间在每批上施加的单步功能。

输入和目标都分别包含批量的输入和目标张量。注意，当我们走向NX生态系统中数据集的统一表示时，输入和目标的格式可能会发生变化。

axon.training.train实现了一个常见的训练循环，初始化训练状态，并通过一些给定数量的时期的培训迭代。它返回序列化的最终培训状态和推理工作负载中的潜在用途。

目前，轴突训练API相当有限;但是，有计划扩展它。在立即的未来，我们计划支持：

此外，我们很乐意探索像分布式培训等更高级的事情。 我们还在寻求通过完全运行本机加速器来提高培训环路的绩效。 轴突仍然很年轻，在它准备发布之前要做多少工作。 你可能会遇到尖锐的边缘，错误和令人困惑的错误。 通过尝试API，我们会更好地热爱您的帮助，提高您找到的任何问题，提供有关我们任何开放问题的任何问题，或者为项目提供反馈。 最后，如果我没有承认一些作为轴突的灵感的图书馆，我会被遗漏：