elixir-nx/axon

何を解決するか

Axon は Elixir エコシステム内でのニューラルネットワークの構築と学習を包括的にサポートするフレームワークです。開発者は主言語から離れることなくディープラーニングモデルを実装でき、モデル作成のための高レベルAPIと、トレーニングループを管理するための構造化APIを提供しつつ、数値定義の低レベル制御を維持します。

動作方法

Axon は Nx（数値 Elixir）の上に構築されており、3つの主要なレイヤーで構成されています：

• 関数型API：defn（数値定義）で構成される基盤で、活性化関数、レイヤー、損失関数などの一般的なディープラーニングコンポーネントを提供します。これにより、Nxコンパイラを介してJITまたはAOTコンパイルが可能になり、高速化が実現できます。
• モデル作成API：Elixirのstructとしてモデルを定義するための高レベルインターフェースです。モデル定義と初期化・適用の責任を分離し、Axon.build/2 を使って予測関数と初期化関数を生成します。
• トレーニングAPI：PyTorch Igniteを参考にしたルーチンベースのシステムで、トレーニングループを実装可能。メトリクスやイベントハンドラをループに挿入して監視・制御できます。

対象ユーザー

自アプリケーションにディープラーニングやニューラルネットワークの構築を直接統合したいElixir開発者、また高レベルのモデル作成と低レベルの数値高速化を組み合わせる柔軟性が必要なユーザー。

特徴

• Nx対応：Nxのテンソルとコンパイラ（例：EXLA）と完全に統合され、ハードウェア加速を実現。
• APIの分離：関数型API、モデル作成API、トレーニングAPIは独立して使用可能。
• コンパイル可能：JITおよびAOTコンパイルをサポート。TensorFlow LiteやONNXなどのフォーマットへの変換も可能。
• 拡張性：Axonの関数型APIを独自の数値定義と組み合わせてカスタム構築が可能。