A high-level interface for creating neural network models.

Axon is built entirely on top of Nx numerical definitions, so every neural network can be JIT or AOT compiled using any Nx compiler, or even transformed into high-level neural network formats like TensorFlow Lite and ONNX.

For a more in-depth overview of Axon, refer to the Guides.

Model Creation

All Axon models start with an input layer, optionally specifying the expected shape of the input data:

input = Axon.input("input", shape: {nil, 784})

Notice you can specify some dimensions as nil, indicating that the dimension size will be filled in at model runtime. You can then compose inputs with other layers:

model =
  input
  |> Axon.dense(128, activation: :relu)
  |> Axon.batch_norm()
  |> Axon.dropout(rate: 0.8)
  |> Axon.dense(64)
  |> Axon.tanh()
  |> Axon.dense(10)
  |> Axon.activation(:softmax)

You can inspect the model for a nice summary:

IO.inspect(model)

#Axon<
  inputs: %{"input" => {nil, 784}}
  outputs: "softmax_0"
  nodes: 9
>

Or use the Axon.Display module to see more in-depth summaries:

Axon.Display.as_table(model, Nx.template({1, 784}, :f32)) |> IO.puts

+----------------------------------------------------------------------------------------------------------------+
|                                                     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]        |
+---------------------------------------+-------------+--------------+-------------------+-----------------------+
| relu_0 ( relu["dense_0"] )            | [{1, 128}]  | {1, 128}     |                   |                       |
+---------------------------------------+-------------+--------------+-------------------+-----------------------+
| batch_norm_0 ( batch_norm["relu_0"] ) | [{1, 128}]  | {1, 128}     | epsilon: 1.0e-5   | gamma: f32[128]       |
|                                       |             |              | channel_index: 1  | beta: f32[128]        |
|                                       |             |              | momentum: 0.1     | mean: f32[128]        |
|                                       |             |              |                   | var: f32[128]         |
+---------------------------------------+-------------+--------------+-------------------+-----------------------+
| dropout_0 ( dropout["batch_norm_0"] ) | [{1, 128}]  | {1, 128}     | rate: 0.8         |                       |
+---------------------------------------+-------------+--------------+-------------------+-----------------------+
| dense_1 ( dense["dropout_0"] )        | [{1, 128}]  | {1, 64}      |                   | kernel: f32[128][64]  |
|                                       |             |              |                   | bias: f32[64]         |
+---------------------------------------+-------------+--------------+-------------------+-----------------------+
| tanh_0 ( tanh["dense_1"] )            | [{1, 64}]   | {1, 64}      |                   |                       |
+---------------------------------------+-------------+--------------+-------------------+-----------------------+
| dense_2 ( dense["tanh_0"] )           | [{1, 64}]   | {1, 10}      |                   | kernel: f32[64][10]   |
|                                       |             |              |                   | bias: f32[10]         |
+---------------------------------------+-------------+--------------+-------------------+-----------------------+
| softmax_0 ( softmax["dense_2"] )      | [{1, 10}]   | {1, 10}      |                   |                       |
+---------------------------------------+-------------+--------------+-------------------+-----------------------+

Multiple Inputs

Creating a model with multiple inputs is as easy as declaring an additional input in your Axon graph. Every input layer present in the final Axon graph will be required to be passed as input at the time of model execution.

inp1 = Axon.input("input_0", shape: {nil, 1})
inp2 = Axon.input("input_1", shape: {nil, 1})

# Both inputs will be used
model1 = Axon.add(inp1, inp2)

# Only inp2 will be used
model2 = Axon.add(inp2, inp2)

Axon graphs are immutable, which means composing and manipulating an Axon graph creates an entirely new graph. Additionally, layer names are lazily generated at model execution time. To avoid non-deterministic input orderings and names, Axon requires each input to have a unique binary identifier. You can then reference inputs by name when passing to models at execution time:

inp1 = Axon.input("input_0", shape: {nil, 1})
inp2 = Axon.input("input_1", shape: {nil, 1})

model1 = Axon.add(inp1, inp2)

{init_fn, predict_fn} = Axon.build(model1)

params1 = init_fn.(Nx.template({1, 1}, {:f, 32}), %{})
# Inputs are referenced by name
predict_fn.(params1, %{"input_0" => x, "input_1" => y})

Multiple Outputs

Nx offers robust container support which is extended to Axon. Axon allows you to wrap any valid Nx container in a layer. Containers are most commonly used to structure outputs:

inp1 = Axon.input("input_0", shape: {nil, 1})
inp2 = Axon.input("input_1", shape: {nil, 1})
model = Axon.container(%{foo: inp1, bar: inp2})

Containers can be arbitrarily nested:

inp1 = Axon.input("input_0", shape: {nil, 1})
inp2 = Axon.input("input_1", shape: {nil, 1})
model = Axon.container({%{foo: {inp1, %{bar: inp2}}}})

You can even use custom structs which implement the container protocol:

inp1 = Axon.input("input_0", shape: {nil, 1})
inp2 = Axon.input("input_1", shape: {nil, 1})
model = Axon.container(%MyStruct{foo: inp1, bar: inp2})

Custom Layers

If you find that Axon's built-in layers are insufficient for your needs, you can create your own using the custom layer API. All of Axon's built-in layers (aside from special ones such as input, constant, and container) make use of this same API.

Axon layers are really just placeholders for Nx computations with trainable parameters and possibly state. To define a custom layer, you just need to define a defn implementation:

defn my_layer(x, weight, _opts \\ []) do
  Nx.atan2(x, weight)
end

Notice the only stipulation is that your custom layer implementation must accept at least 1 input and a list of options. At execution time, every layer will be passed a :mode option which can be used to control behavior at training and inference time.

Inputs to your custom layer can be either Axon graph inputs or trainable parameters. You can pass Axon graph inputs as-is to a custom layer. To declare trainable parameters, use Axon.param/3:

weight = Axon.param("weight", param_shape)

To create a custom layer, you "wrap" your implementation and inputs into a layer using Axon.layer. You'll notice the API mirrors Elixir's apply:

def atan2_layer(%Axon{} = input) do
  weight = Axon.param("weight", param_shape)
  Axon.layer(&my_layer/3, [input, weight])
end

Model Execution

Under the hood, Axon models are represented as Elixir structs. You can initialize and apply models by building or compiling them with Axon.build/2 or Axon.compile/4 and then calling the produced initialization and predict functions:

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

params = init_fn.(Nx.template({1, 1}, {:f, 32}), %{})
predict_fn.(params, inputs)

You may either set the default JIT compiler or backend globally, or pass a specific compiler to Axon.build/2:

EXLA.set_as_nx_default([:tpu, :cuda, :rocm, :host])

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

params = init_fn.(Nx.template({1, 1}, {:f, 32}), %{})
predict_fn.(params, inputs)

predict_fn by default runs in inference mode, which performs certain optimizations and removes layers such as dropout layers. If constructing a training step using Axon.predict/4 or Axon.build/2, be sure to specify mode: :train.

Model Training

Combining the Axon model creation API with the optimization and training APIs, you can create and train neural networks with ease:

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

IO.inspect model

model_state =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(learning_rate: 0.005))
  |> Axon.Loop.run(train_data, epochs: 10, compiler: EXLA)

See Polaris.Updates and Axon.Loop for a more in-depth treatment of model optimization and model training.

Using with Nx.Serving

When deploying an Axon model to production, you usually want to batch multiple prediction requests and run the inference for all of them at once. Conveniently, Nx already has an abstraction for this task in the form of Nx.Serving. Here's how you could define a serving for an Axon model:

def build_serving() do
  # Configuration
  batch_size = 4
  defn_options = [compiler: EXLA]

  Nx.Serving.new(
    # This function runs on the serving startup
    fn ->
      # Build the Axon model and load params (usually from file)
      model = build_model()
      params = load_params()

      # Build the prediction defn function
      {_init_fun, predict_fun} = Axon.build(model)

      inputs_template = %{"pixel_values" => Nx.template({batch_size, 224, 224, 3}, :f32)}
      template_args = [Nx.to_template(params), inputs_template]

      # Compile the prediction function upfront for the configured batch_size
      predict_fun = Nx.Defn.compile(predict_fun, template_args, defn_options)

      # The returned function is called for every accumulated batch
      fn inputs ->
        inputs = Nx.Batch.pad(inputs, batch_size - inputs.size)
        predict_fun.(params, inputs)
      end
    end,
    batch_size: batch_size
  )
end

Then you would start the serving server as part of your application's supervision tree:

children = [
  ...,
  {Nx.Serving, serving: build_serving(), name: MyApp.Serving, batch_timeout: 100}
]

With that in place, you can now ask serving for predictions all across your application (controllers, live views, async jobs, etc.). Having a tensor input you would do:

inputs = %{"pixel_values" => ...}
batch = Nx.Batch.concatenate([inputs])
result = Nx.Serving.batched_run(MyApp.Serving, batch)

Usually you also want to do pre/post-processing of the model input/output. You could make those preparations directly before/after Nx.Serving.batched_run/2, however you can also make use of Nx.Serving.client_preprocessing/2 and Nx.Serving.client_postprocessing/2 to encapsulate that logic as part of the serving.