defmodule Axon do
@moduledoc """
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](https://github.com/elixir-nx/axon_onnx).
## Model Creation
All Axon models start with an input layer, specifying the
expected input shape of the training data:
input = Axon.input({nil, 784}, "input")
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)
---------------------------------------------------------------------------------------------------------
Model
=========================================================================================================
Layer Shape Policy Parameters Parameters Memory
=========================================================================================================
input ( input ) {nil, 784} p=f32 c=f32 o=f32 0 0 bytes
dense_0 ( dense["input"] ) {nil, 128} p=f32 c=f32 o=f32 100480 401920 bytes
relu_0 ( relu["dense_0"] ) {nil, 128} p=f32 c=f32 o=f32 0 0 bytes
batch_norm_0 ( batch_norm["relu_0"] ) {nil, 128} p=f32 c=f32 o=f32 512 2048 bytes
dropout_0 ( dropout["batch_norm_0"] ) {nil, 128} p=f32 c=f32 o=f32 0 0 bytes
dense_1 ( dense["dropout_0"] ) {nil, 64} p=f32 c=f32 o=f32 8256 33024 bytes
tanh_0 ( tanh["dense_1"] ) {nil, 64} p=f32 c=f32 o=f32 0 0 bytes
dense_2 ( dense["tanh_0"] ) {nil, 10} p=f32 c=f32 o=f32 650 2600 bytes
softmax_0 ( softmax["dense_2"] ) {nil, 10} p=f32 c=f32 o=f32 0 0 bytes
---------------------------------------------------------------------------------------------------------
Total Parameters: 109898
Total Parameters Memory: 439592 bytes
Inputs: %{"input" => {nil, 784}}
### 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({nil, 1}, "input_0")
inp2 = Axon.input({nil, 1}, "input_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({nil, 1}, "input_0")
inp2 = Axon.input({nil, 1}, "input_1")
model1 = Axon.add(inp1, inp2)
params1 = Axon.init(model1)
# Inputs are referenced by name
Axon.predict(model1, params1, %{"input_0" => x, "input_1" => y})
### Multiple Outputs
Nx offers robust [container](https://hexdocs.pm/nx/Nx.Container.html) 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({nil, 1}, "input_0")
inp2 = Axon.input({nil, 1}, "input_1")
model = Axon.container(%{foo: inp1, bar: inp2})
Containers can be arbitrarily nested:
inp1 = Axon.input({nil, 1}, "input_0")
inp2 = Axon.input({nil, 1}, "input_1")
model = Axon.container({%{foo: {inp1, %{bar: inp2}}}})
You can even use custom structs which implement the container protocol:
inp1 = Axon.input({nil, 1}, "input_0")
inp2 = Axon.input({nil, 1}, "input_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(input_shape, "weight")
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{output_shape: shape} = input) do
weight = Axon.param(input_shape, "weight")
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 using the macros `Axon.init/3` and
`Axon.predict/4`:
params = Axon.init(model, compiler: EXLA)
Axon.predict(model, params, inputs, compiler: EXLA, mode: :train)
It is suggested that you set compiler options globally rather than pass
them as options to execution macros:
EXLA.set_as_nx_default([:tpu, :cuda, :rocm, :host])
params = Axon.init(model)
Axon.predict(model, params, inputs, mode: :train)
`Axon.predict/4` 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`, 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({nil, 784}, "input_0")
|> 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, Axon.Optimizers.adamw(0.005))
|> Axon.Loop.run(train_data, epochs: 10, compiler: EXLA)
See `Axon.Updates` and `Axon.Loop` for a more in-depth treatment of
model optimization and model training.
"""
alias __MODULE__, as: Axon
alias Axon.Parameter
# Axon serialization version
@file_version 1
@type t :: %__MODULE__{}
defstruct [
:id,
:name,
:output_shape,
:parent,
:parameters,
:args,
:op,
:policy,
:hooks,
:opts,
:op_name
]
@doc """
Custom Axon layer with given inputs.
Inputs may be other Axon layers or trainable parameters created
with `Axon.param`. At inference time, `op` will be applied with
inputs in specified order and an additional `opts` parameter which
specifies inference options. All options passed to layer are forwarded
to inference function except:
* `:shape` - specify layer output shape to bypass shape inference.
* `:name` - layer name.
* `:op_name` - layer operation for inspection and building parameter
map.
Note this means your layer should not use these as input options,
as they will always be dropped during inference compilation.
Axon's compiler will additionally forward the following options to
every layer at inference time:
* `:mode` - `:inference` or `:train`. To control layer behavior
based on inference or train time.
`op` is a function of the form:
fun = fn input, weight, bias, _opts ->
input * weight + bias
end
"""
@doc type: :special
def layer(op, inputs, opts \\ []) when (is_atom(op) or is_function(op)) and is_list(inputs) do
{inputs, params, args, input_shapes} = split_inputs(op, inputs)
inputs = Enum.reverse(inputs)
params = Enum.reverse(params)
args = Enum.reverse(args)
input_shapes = Enum.reverse(input_shapes)
{name, opts} = Keyword.pop(opts, :name)
{shape, opts} = Keyword.pop(opts, :shape)
{op_name, opts} = Keyword.pop(opts, :op_name, :custom)
{id, name} = unique_identifiers(op_name, name)
output_shape =
if shape do
shape
else
infer_shape(input_shapes, op, opts)
end
%Axon{
id: id,
name: name,
output_shape: output_shape,
parent: inputs,
parameters: params,
args: args,
op: op,
policy: Axon.MixedPrecision.create_policy(),
hooks: [],
opts: opts,
op_name: op_name
}
end
defp split_inputs(:container, [container] = inputs) do
input_shapes = deep_new(container, fn %Axon{output_shape: shape} -> shape end)
args = [:layer]
params = []
{inputs, params, args, [input_shapes]}
end
defp split_inputs(_op, inputs) do
Enum.reduce(inputs, {[], [], [], []}, fn
%Axon{output_shape: shape} = layer, {layers, params, args, shapes} ->
{[layer | layers], params, [:layer | args], [shape | shapes]}
%Parameter{shape: shape} = param, {layers, params, args, shapes} ->
{layers, [param | params], [:parameter | args], [shape | shapes]}
invalid, _ ->
raise ArgumentError, "invalid input given to layer: #{inspect(invalid)}"
end)
end
defp infer_shape(input_shapes, fun, opts) do
{inputs, indices} =
Enum.reduce(input_shapes, {[], []}, fn shape, {input_shapes, indices} ->
{template, template_indices} = template_shape(shape)
{[template | input_shapes], [template_indices | indices]}
end)
inputs = Enum.reverse(inputs)
opts = Keyword.put(opts, :mode, :inference)
wrapper_fun = fn tensors ->
tensors = Tuple.to_list(tensors)
apply(fun, tensors ++ [opts])
end
expr = Nx.Defn.jit(wrapper_fun, [List.to_tuple(inputs)], compiler: Axon.Defn)
indices = Enum.map(indices, &MapSet.new/1)
indices_that_are_1 =
expr.shape
|> Tuple.to_list()
|> Enum.with_index()
|> Enum.filter(fn {x, _} -> x == 1 end)
|> Enum.map(fn {_, i} -> i end)
|> MapSet.new()
indices_to_make_nil =
case indices do
[] ->
[]
indices ->
indices
|> Enum.reduce(MapSet.new(), &MapSet.union/2)
|> MapSet.intersection(indices_that_are_1)
|> Enum.to_list()
end
Enum.reduce(indices_to_make_nil, expr.shape, fn i, shape ->
put_elem(shape, i, nil)
end)
end
defp template_shape(shape) when is_map(shape) do
Nx.Container.traverse(shape, [], &recur_template_shape/2)
end
defp template_shape(shape) do
if tuple_size(shape) == 0 do
{Nx.template({}, {:f, 32}), []}
else
first_elem = elem(shape, 0)
if is_integer(first_elem) or is_nil(first_elem) do
{shape, template_indices} = Axon.Shape.replace_nil(shape)
template = Nx.template(shape, {:f, 32})
{template, List.wrap(template_indices)}
else
Nx.Container.traverse(shape, [], &recur_template_shape/2)
end
end
end
defp recur_template_shape(shape, indices) do
case shape do
shape when is_map(shape) ->
{template, template_indices} = template_shape(shape)
{template, indices ++ template_indices}
shape when is_tuple(shape) ->
{template, template_indices} = template_shape(shape)
{template, indices ++ template_indices}
end
end
@doc """
Trainable Axon parameter used to create custom layers.
Parameters are specified in usages of `Axon.layer` and will
be automatically initialized and used in subsequent applications
of Axon models.
Parameters *must* be specified in order of their usage.
## Options
* `:initializer` - parameter initializer. Defaults to `:glorot_uniform`.
"""
def param(name, shape, opts \\ []) when is_binary(name) and is_tuple(shape) do
opts = Keyword.validate!(opts, initializer: :glorot_uniform)
initializer = opts[:initializer]
validate_initializer!(initializer)
id = System.unique_integer([:positive, :monotonic])
%Axon.Parameter{
id: id,
name: name,
shape: shape,
initializer: initializer
}
end
@doc """
Adds an input layer to the network.
Input layers specify a model's inputs. Input layers are
always the root layers of the neural network.
You must specify the input layers name, which will be used
to uniquely identify it in the case of multiple inputs.
"""
@doc type: :special
def input(input_shape, name) when is_tuple(input_shape) and is_binary(name) do
output_shape = Axon.Shape.input(input_shape)
layer(:input, [], name: name, shape: output_shape, op_name: :input)
end
@doc """
Adds a constant layer to the network.
Constant layers encapsulate Nx tensors in an Axon layer for ease
of use with other Axon layers. They can be used interchangeably
with other Axon layers:
inp = Axon.input({nil, 32}, "input")
my_constant = Axon.constant(Nx.iota({1, 32}))
model = Axon.add(inp, my_constant)
Constant layers will be cast according to the mixed precision policy.
If it's important for your constant to retain it's type during
the computation, you will need to set the mixed precision policy to
ignore constant layers.
## Options
* `:name` - layer name.
"""
def constant(tensor, opts \\ [])
@doc type: :special
def constant(%Nx.Tensor{shape: output_shape} = tensor, opts) do
opts = Keyword.validate!(opts, [:name])
layer(:constant, [], name: opts[:name], value: tensor, shape: output_shape, op_name: :constant)
end
def constant(value, _) do
raise ArgumentError,
"value passed to constant must be an Nx tensor" <>
" but got #{inspect(value)}, if you are passing" <>
" a number, wrap it with a call to Nx.tensor/2"
end
@doc """
Adds a container layer to the network.
In certain cases you may want your model to have multiple
outputs. In order to make this work, you must "join" the
outputs into an Axon layer using this function for use in
initialization and inference later on.
The given container can be any valid Axon Nx container.
## Options
* `:name` - layer name.
## Examples
iex> inp1 = Axon.input({nil, 1}, "input_0")
iex> inp2 = Axon.input({nil, 2}, "input_1")
iex> model = Axon.container(%{a: inp1, b: inp2})
iex> %{a: a, b: b} = Axon.predict(model, %{}, %{
...> "input_0" => Nx.tensor([[1.0]]),
...> "input_1" => Nx.tensor([[1.0, 2.0]])
...> })
iex> a
#Nx.Tensor<
f32[1][1]
[
[1.0]
]
>
iex> b
#Nx.Tensor<
f32[1][2]
[
[1.0, 2.0]
]
>
"""
@doc type: :special
def container(container, opts \\ []) do
opts = Keyword.validate!(opts, [:name])
output_shape =
deep_new(container, fn %Axon{output_shape: shape} ->
shape
end)
layer(:container, [container], name: opts[:name], shape: output_shape, op_name: :container)
end
# TODO: This should not be duplicated
defp deep_new(map, fun) do
{cont, :ok} = Nx.Container.traverse(map, :ok, &recur_traverse(&1, &2, fun))
cont
end
defp recur_traverse(item, :ok, fun) do
case item do
%Axon{} = t ->
{fun.(t), :ok}
%{axon: :axon} = t ->
{fun.(t), :ok}
container ->
{deep_new(container, fun), :ok}
end
end
@doc """
Wraps an Axon model into a namespace.
A namespace is a part of an Axon model which is meant to
be a self-contained collection of Axon layers. Namespaces
are guaranteed to always generate with the same internal
layer names and can be re-used universally across models.
Namespaces are most useful for containing large collections
of layers and offering a straightforward means for accessing
the parameters of individual model components. A common application
of namespaces is to use them in with a pre-trained model for
fine-tuning:
{base, resnet_params} = resnet()
base = base |> Axon.namespace("resnet")
model = base |> Axon.dense(1)
Axon.init(model, %{"resnset" => resnet_params})
Notice you can use `Axon.init` in conjunction with namespaces
to specify which portion of a model you'd like to initialize
from a fixed starting point.
Namespaces have fixed names, which means it's easy to run into namespace
collisions. Re-using namespaces, re-using inner parts of a namespace,
and attempting to share layers between namespaces are still sharp
edges in namespace usage.
"""
def namespace(%Axon{output_shape: shape} = axon, name) when is_binary(name) do
layer(:namespace, [axon], name: name, shape: shape)
end
@doc """
Adds a dense layer to the network.
The dense layer implements:
output = activation(dot(input, kernel) + bias)
where `activation` is given by the `:activation` option and both
`kernel` and `bias` are layer parameters. `units` specifies the
number of output units.
Compiles to `Axon.Layers.dense/4`.
## Options
* `:name` - layer name.
* `:kernel_initializer` - initializer for `kernel` weights.
Defaults to `:glorot_uniform`.
* `:bias_initializer` - initializer for `bias` weights. Defaults
to `:zeros`.
* `:activation` - element-wise activation function.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`.
"""
@doc type: :linear
def dense(%Axon{output_shape: parent_shape} = x, units, opts \\ [])
when is_integer(units) and units > 0 do
opts =
Keyword.validate!(opts, [
:name,
:activation,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros,
use_bias: true
])
kernel_shape = Axon.Shape.dense_kernel(parent_shape, units)
bias_shape = Axon.Shape.dense_bias(parent_shape, units)
output_shape = Axon.Shape.dense(parent_shape, units)
kernel = param("kernel", kernel_shape, initializer: opts[:kernel_initializer])
{inputs, op} =
if opts[:use_bias] do
bias = param("bias", bias_shape, initializer: opts[:bias_initializer])
{[x, kernel, bias], :dense}
else
{[x, kernel], &Axon.Layers.dense(&1, &2, 0, &3)}
end
node = layer(op, inputs, name: opts[:name], shape: output_shape, op_name: :dense)
if activation = opts[:activation] do
activation(node, activation)
else
node
end
end
@doc """
Adds a bilinear layer to the network.
The bilinear layer implements:
output = activation(dot(dot(input1, kernel), input2) + bias)
where `activation` is given by the `:activation` option and both
`kernel` and `bias` are layer parameters. `units` specifies the
number of output units.
All dimensions but the last of `input1` and `input2` must match. The
batch sizes of both inputs must also match or at least one must be `nil`.
Inferred output batch size coerces to the strictest input batch size.
Compiles to `Axon.Layers.bilinear/5`.
## Options
* `:name` - layer name.
* `:kernel_initializer` - initializer for `kernel` weights.
Defaults to `:glorot_uniform`.
* `:bias_initializer` - initializer for `bias` weights. Defaults
to `:zeros`.
* `:activation` - element-wise activation function.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`.
"""
@doc type: :linear
def bilinear(
%Axon{output_shape: parent1_shape} = input1,
%Axon{output_shape: parent2_shape} = input2,
units,
opts \\ []
)
when is_integer(units) and units > 0 do
opts =
Keyword.validate!(opts, [
:name,
:activation,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros,
use_bias: true
])
kernel_shape = Axon.Shape.bilinear_kernel(parent1_shape, parent2_shape, units)
bias_shape = Axon.Shape.bilinear_bias(parent1_shape, parent2_shape, units)
output_shape = Axon.Shape.bilinear(parent1_shape, parent2_shape, units)
kernel = param("kernel", kernel_shape, initializer: opts[:kernel_initializer])
{inputs, op} =
if opts[:use_bias] do
bias = param("bias", bias_shape, initializer: opts[:bias_initializer])
{[input1, input2, kernel, bias], :bilinear}
else
{[input1, input2, kernel], &Axon.Layers.bilinear(&1, &2, &3, 0, &4)}
end
node = layer(op, inputs, name: opts[:name], shape: output_shape, op_name: :bilinear)
if activation = opts[:activation] do
activation(node, activation)
else
node
end
end
@doc """
Adds a convolution layer to the network.
The convolution layer implements a general dimensional
convolutional layer - which convolves a kernel over the input
to produce an output.
Compiles to `Axon.Layers.conv/4`.
## Options
* `:name` - layer name.
* `:kernel_initializer` - initializer for `kernel` weights.
Defaults to `:glorot_uniform`.
* `:bias_initializer` - initializer for `bias` weights. Defaults
to `:zeros`
* `:activation` - element-wise activation function.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`
* `:kernel_size` - size of the kernel spatial dimensions. Defaults
to `1`.
* `:strides` - stride during convolution. Defaults to `1`.
* `:padding` - padding to the spatial dimensions of the input.
Defaults to `:valid`.
* `:input_dilation` - dilation to apply to input. Defaults to `1`.
* `:kernel_dilation` - dilation to apply to kernel. Defaults to `1`.
* `:feature_group_size` - feature group size for convolution. Defaults
to `1`.
* `:channels` - channels location. One of `:first` or `:last`.
Defaults to `:first`.
"""
@doc type: :convolution
def conv(%Axon{output_shape: parent_shape} = x, units, opts \\ [])
when is_integer(units) and units > 0 do
opts =
Keyword.validate!(opts, [
:name,
:activation,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros,
use_bias: true,
kernel_size: 1,
strides: 1,
padding: :valid,
input_dilation: 1,
kernel_dilation: 1,
channels: :first,
feature_group_size: 1
])
kernel_size = opts[:kernel_size]
strides = opts[:strides]
padding = opts[:padding]
input_dilation = opts[:input_dilation]
kernel_dilation = opts[:kernel_dilation]
channels = opts[:channels]
feature_group_size = opts[:feature_group_size]
inner_rank = Nx.rank(parent_shape) - 2
kernel_size = tuple_or_duplicate(:kernel_size, kernel_size, inner_rank)
strides = list_or_duplicate(:strides, strides, inner_rank)
input_dilation = list_or_duplicate(:input_dilation, input_dilation, inner_rank)
kernel_dilation = list_or_duplicate(:kernel_dilation, kernel_dilation, inner_rank)
kernel_shape = Axon.Shape.conv_kernel(parent_shape, units, kernel_size, channels)
bias_shape = Axon.Shape.conv_bias(parent_shape, units, kernel_size, channels)
output_shape =
Axon.Shape.conv(
parent_shape,
kernel_shape,
strides,
padding,
input_dilation,
kernel_dilation,
channels,
feature_group_size
)
kernel = param("kernel", kernel_shape, initializer: opts[:kernel_initializer])
{inputs, op} =
if opts[:use_bias] do
bias = param("bias", bias_shape, initializer: opts[:bias_initializer])
{[x, kernel, bias], :conv}
else
{[x, kernel], &Axon.Layers.conv(&1, &2, 0, &3)}
end
node =
layer(op, inputs,
name: opts[:name],
strides: strides,
padding: padding,
input_dilation: input_dilation,
kernel_dilation: kernel_dilation,
feature_group_size: feature_group_size,
channels: channels,
shape: output_shape,
op_name: :conv
)
if activation = opts[:activation] do
activation(node, activation)
else
node
end
end
@doc """
Adds a transposed convolution layer to the network.
The transposed convolution layer is sometimes referred to as a
fractionally strided convolution or (incorrectly) as a deconvolution.
Compiles to `Axon.Layers.conv_transpose/4`.
## Options
* `:name` - layer name.
* `:kernel_initializer` - initializer for `kernel` weights.
Defaults to `:glorot_uniform`.
* `:bias_initializer` - initializer for `bias` weights. Defaults
to `:zeros`
* `:activation` - element-wise activation function.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`
* `:kernel_size` - size of the kernel spatial dimensions. Defaults
to `1`.
* `:strides` - stride during convolution. Defaults to `1`.
* `:padding` - padding to the spatial dimensions of the input.
Defaults to `:valid`.
* `:kernel_dilation` - dilation to apply to kernel. Defaults to `1`.
* `:channels` - channels location. One of `:first` or `:last`.
Defaults to `:first`.
"""
@doc type: :convolution
def conv_transpose(%Axon{output_shape: parent_shape} = x, units, opts \\ []) do
opts =
Keyword.validate!(opts, [
:name,
:activation,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros,
use_bias: true,
kernel_size: 1,
strides: 1,
padding: :valid,
kernel_dilation: 1,
channels: :first
])
kernel_size = opts[:kernel_size]
strides = opts[:strides]
padding = opts[:padding]
kernel_dilation = opts[:kernel_dilation]
channels = opts[:channels]
inner_rank = Nx.rank(parent_shape) - 2
kernel_size = tuple_or_duplicate(:kernel_size, kernel_size, inner_rank)
strides = list_or_duplicate(:strides, strides, inner_rank)
kernel_dilation = list_or_duplicate(:kernel_dilation, kernel_dilation, inner_rank)
kernel_shape = Axon.Shape.conv_kernel(parent_shape, units, kernel_size, channels)
bias_shape = Axon.Shape.conv_bias(parent_shape, units, kernel_size, channels)
kernel = param("kernel", kernel_shape, initializer: opts[:kernel_initializer])
{inputs, op} =
if opts[:use_bias] do
bias = param("bias", bias_shape, initializer: opts[:bias_initializer])
{[x, kernel, bias], :conv_transpose}
else
{[x, kernel], &Axon.Layers.conv_transpose(&1, &2, 0, &3)}
end
output_shape =
Axon.Shape.conv_transpose(
parent_shape,
kernel_shape,
strides,
padding,
kernel_dilation,
channels
)
node =
layer(op, inputs,
name: opts[:name],
strides: strides,
padding: padding,
kernel_dilation: kernel_dilation,
channels: channels,
shape: output_shape,
op_name: :conv_transpose
)
if activation = opts[:activation] do
activation(node, activation)
else
node
end
end
@doc """
Adds a depthwise convolution layer to the network.
The depthwise convolution layer implements a general
dimensional depthwise convolution - which is a convolution
where the feature group size is equal to the number of
input channels.
Channel multiplier grows the input channels by the given
factor. An input factor of 1 means the output channels
are the same as the input channels.
Compiles to `Axon.Layers.depthwise_conv/4`.
## Options
* `:name` - layer name.
* `:kernel_initializer` - initializer for `kernel` weights.
Defaults to `:glorot_uniform`.
* `:bias_initializer` - initializer for `bias` weights. Defaults
to `:zeros`
* `:activation` - element-wise activation function.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`
* `:kernel_size` - size of the kernel spatial dimensions. Defaults
to `1`.
* `:strides` - stride during convolution. Defaults to `1`.
* `:padding` - padding to the spatial dimensions of the input.
Defaults to `:valid`.
* `:input_dilation` - dilation to apply to input. Defaults to `1`.
* `:kernel_dilation` - dilation to apply to kernel. Defaults to `1`.
* `:channels` - channels location. One of `:first` or `:last`.
Defaults to `:first`.
"""
@doc type: :convolution
def depthwise_conv(%Axon{output_shape: parent_shape} = x, channel_multiplier, opts \\ [])
when is_integer(channel_multiplier) and channel_multiplier >= 1 do
opts =
Keyword.validate!(opts, [
:name,
:activation,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros,
use_bias: true,
kernel_size: 1,
strides: 1,
padding: :valid,
input_dilation: 1,
kernel_dilation: 1,
channels: :first
])
kernel_size = opts[:kernel_size]
strides = opts[:strides]
padding = opts[:padding]
input_dilation = opts[:input_dilation]
kernel_dilation = opts[:kernel_dilation]
channels = opts[:channels]
inner_rank = Nx.rank(parent_shape) - 2
kernel_size = tuple_or_duplicate(:kernel_size, kernel_size, inner_rank)
strides = list_or_duplicate(:strides, strides, inner_rank)
input_dilation = list_or_duplicate(:input_dilation, input_dilation, inner_rank)
kernel_dilation = list_or_duplicate(:kernel_dilation, kernel_dilation, inner_rank)
kernel_shape =
Axon.Shape.depthwise_conv_kernel(parent_shape, channel_multiplier, kernel_size, channels)
bias_shape =
Axon.Shape.depthwise_conv_bias(parent_shape, channel_multiplier, kernel_size, channels)
output_shape =
Axon.Shape.depthwise_conv(
parent_shape,
kernel_shape,
strides,
padding,
input_dilation,
kernel_dilation,
channels
)
kernel = param("kernel", kernel_shape, initializer: opts[:kernel_initializer])
{inputs, op} =
if opts[:use_bias] do
bias = param("bias", bias_shape, initializer: opts[:bias_initializer])
{[x, kernel, bias], :depthwise_conv}
else
{[x, kernel], &Axon.Layers.depthwise_conv(&1, &2, 0, &3)}
end
node =
layer(op, inputs,
name: opts[:name],
strides: strides,
padding: padding,
input_dilation: input_dilation,
kernel_dilation: kernel_dilation,
channels: channels,
shape: output_shape,
op_name: :depthwise_conv
)
if activation = opts[:activation] do
activation(node, activation)
else
node
end
end
@doc """
Adds a depthwise separable 2-dimensional convolution to the
network.
Depthwise separable convolutions break the kernel into kernels
for each dimension of the input and perform a depthwise conv
over the input with each kernel.
Compiles to `Axon.Layers.separable_conv2d/6`.
## Options
* `:name` - layer name.
* `:kernel_initializer` - initializer for `kernel` weights.
Defaults to `:glorot_uniform`.
* `:bias_initializer` - initializer for `bias` weights. Defaults
to `:zeros`
* `:activation` - element-wise activation function.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`
* `:kernel_size` - size of the kernel spatial dimensions. Defaults
to `1`.
* `:strides` - stride during convolution. Defaults to `1`.
* `:padding` - padding to the spatial dimensions of the input.
Defaults to `:valid`.
* `:input_dilation` - dilation to apply to input. Defaults to `1`.
* `:kernel_dilation` - dilation to apply to kernel. Defaults to `1`.
* `:channels` - channels location. One of `:first` or `:last`.
Defaults to `:first`.
"""
@doc type: :convolution
def separable_conv2d(%Axon{output_shape: parent_shape} = x, channel_multiplier, opts \\ [])
when is_integer(channel_multiplier) and channel_multiplier >= 1 do
opts =
Keyword.validate!(opts, [
:name,
:activation,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros,
use_bias: true,
kernel_size: 1,
strides: 1,
padding: :valid,
input_dilation: 1,
kernel_dilation: 1,
channels: :first
])
kernel_size = opts[:kernel_size]
strides = opts[:strides]
padding = opts[:padding]
input_dilation = opts[:input_dilation]
kernel_dilation = opts[:kernel_dilation]
channels = opts[:channels]
inner_rank = Nx.rank(parent_shape) - 2
kernel_size = tuple_or_duplicate(:kernel_size, kernel_size, inner_rank)
strides = list_or_duplicate(:strides, strides, inner_rank)
input_dilation = list_or_duplicate(:input_dilation, input_dilation, inner_rank)
kernel_dilation = list_or_duplicate(:kernel_dilation, kernel_dilation, inner_rank)
k1_shape =
Axon.Shape.separable_conv2d_kernel(
parent_shape,
channel_multiplier,
kernel_size,
1,
channels
)
k2_shape =
Axon.Shape.separable_conv2d_kernel(
parent_shape,
channel_multiplier,
kernel_size,
2,
channels
)
b1_shape =
Axon.Shape.separable_conv2d_bias(parent_shape, channel_multiplier, kernel_size, channels)
b2_shape =
Axon.Shape.separable_conv2d_bias(parent_shape, channel_multiplier, kernel_size, channels)
output_shape =
Axon.Shape.depthwise_conv(
parent_shape,
Axon.Shape.depthwise_conv_kernel(parent_shape, channel_multiplier, kernel_size, channels),
strides,
padding,
input_dilation,
kernel_dilation,
channels
)
kernel_initializer = opts[:kernel_initializer]
k1 = param("kernel_1", k1_shape, initializer: kernel_initializer)
k2 = param("kernel_2", k2_shape, initializer: kernel_initializer)
{inputs, op} =
if opts[:use_bias] do
bias_initializer = opts[:bias_initializer]
b1 = param("bias_1", b1_shape, initializer: bias_initializer)
b2 = param("bias_2", b2_shape, initializer: bias_initializer)
{[x, k1, b1, k2, b2], :separable_conv2d}
else
{[x, k1, k2], &Axon.Layers.separable_conv2d(&1, &2, 0, &3, 0, &4)}
end
node =
layer(
op,
inputs,
name: opts[:name],
strides: strides,
padding: padding,
input_dilation: input_dilation,
kernel_dilation: kernel_dilation,
channels: channels,
shape: output_shape,
op_name: :separable_conv2d
)
if activation = opts[:activation] do
activation(node, activation)
else
node
end
end
@doc """
Adds a depthwise separable 3-dimensional convolution to the
network.
Depthwise separable convolutions break the kernel into kernels
for each dimension of the input and perform a depthwise conv
over the input with each kernel.
Compiles to `Axon.Layers.separable_conv3d/8`.
## Options
* `:name` - layer name.
* `:kernel_initializer` - initializer for `kernel` weights.
Defaults to `:glorot_uniform`.
* `:bias_initializer` - initializer for `bias` weights. Defaults
to `:zeros`
* `:activation` - element-wise activation function.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`
* `:kernel_size` - size of the kernel spatial dimensions. Defaults
to `1`.
* `:strides` - stride during convolution. Defaults to `1`.
* `:padding` - padding to the spatial dimensions of the input.
Defaults to `:valid`.
* `:input_dilation` - dilation to apply to input. Defaults to `1`.
* `:kernel_dilation` - dilation to apply to kernel. Defaults to `1`.
* `:channels` - channels location. One of `:first` or `:last`.
Defaults to `:first`.
"""
@doc type: :convolution
def separable_conv3d(%Axon{output_shape: parent_shape} = x, channel_multiplier, opts \\ [])
when is_integer(channel_multiplier) and channel_multiplier >= 1 do
opts =
Keyword.validate!(opts, [
:name,
:activation,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros,
use_bias: true,
kernel_size: 1,
strides: 1,
padding: :valid,
input_dilation: 1,
kernel_dilation: 1,
channels: :first
])
kernel_size = opts[:kernel_size]
strides = opts[:strides]
padding = opts[:padding]
input_dilation = opts[:input_dilation]
kernel_dilation = opts[:kernel_dilation]
channels = opts[:channels]
inner_rank = Nx.rank(parent_shape) - 2
kernel_size = tuple_or_duplicate(:kernel_size, kernel_size, inner_rank)
strides = list_or_duplicate(:strides, strides, inner_rank)
input_dilation = list_or_duplicate(:input_dilation, input_dilation, inner_rank)
kernel_dilation = list_or_duplicate(:kernel_dilation, kernel_dilation, inner_rank)
k1_shape =
Axon.Shape.separable_conv3d_kernel(
parent_shape,
channel_multiplier,
kernel_size,
1,
channels
)
k2_shape =
Axon.Shape.separable_conv3d_kernel(
parent_shape,
channel_multiplier,
kernel_size,
2,
channels
)
k3_shape =
Axon.Shape.separable_conv3d_kernel(
parent_shape,
channel_multiplier,
kernel_size,
3,
channels
)
b1_shape =
Axon.Shape.separable_conv3d_bias(parent_shape, channel_multiplier, kernel_size, channels)
b2_shape =
Axon.Shape.separable_conv3d_bias(parent_shape, channel_multiplier, kernel_size, channels)
b3_shape =
Axon.Shape.separable_conv3d_bias(parent_shape, channel_multiplier, kernel_size, channels)
output_shape =
Axon.Shape.depthwise_conv(
parent_shape,
Axon.Shape.depthwise_conv_kernel(parent_shape, channel_multiplier, kernel_size, channels),
strides,
padding,
input_dilation,
kernel_dilation,
channels
)
kernel_initializer = opts[:kernel_initializer]
k1 = param("kernel_1", k1_shape, initializer: kernel_initializer)
k2 = param("kernel_2", k2_shape, initializer: kernel_initializer)
k3 = param("kernel_3", k3_shape, initializer: kernel_initializer)
{inputs, op} =
if opts[:use_bias] do
bias_initializer = opts[:bias_initializer]
b1 = param("bias_1", b1_shape, initializer: bias_initializer)
b2 = param("bias_2", b2_shape, initializer: bias_initializer)
b3 = param("bias_3", b3_shape, initializer: bias_initializer)
{[x, k1, b1, k2, b2, k3, b3], :separable_conv3d}
else
{[x, k1, k2, k3], &Axon.Layers.separable_conv3d(&1, &2, 0, &3, 0, &4, 0, &5)}
end
node =
layer(
op,
inputs,
name: opts[:name],
strides: strides,
padding: padding,
input_dilation: input_dilation,
kernel_dilation: kernel_dilation,
channels: channels,
shape: output_shape,
op_name: :separable_conv3d
)
if activation = opts[:activation] do
activation(node, activation)
else
node
end
end
@activation_layers [
{:celu, "Continuously-differentiable exponential linear unit", "a"},
{:elu, "Exponential linear unit", "an"},
{:exp, "Exponential", "an"},
{:gelu, "Gaussian error linear unit", "a"},
{:hard_sigmoid, "Hard sigmoid", "a"},
{:hard_silu, "Hard sigmoid weighted linear unit", "a"},
{:hard_tanh, "Hard hyperbolic tangent", "a"},
{:leaky_relu, "Leaky rectified linear unit", "a"},
{:linear, "Linear", "a"},
{:log_sigmoid, "Log-sigmoid", "a"},
{:log_softmax, "Log-softmax", "a"},
{:mish, "Mish", "a"},
{:relu, "Rectified linear unit", "a"},
{:relu6, "Rectified linear unit 6", "a"},
{:sigmoid, "Sigmoid", "a"},
{:silu, "Sigmoid weighted linear unit", "a"},
{:selu, "Scaled exponential linear unit", "a"},
{:softmax, "Softmax", "a"},
{:softplus, "Softplus", "a"},
{:softsign, "Softsign", "a"},
{:tanh, "Hyperbolic tangent", "a"}
]
@doc """
Adds an activation layer to the network.
Activation layers are element-wise functions typically called
after the output of another layer.
## Options
* `:name` - layer name.
"""
@doc type: :activation
def activation(x, activation, opts \\ [])
def activation(%Axon{output_shape: shape} = x, activation, opts) when is_atom(activation) do
opts = [shape: shape, op_name: activation] ++ opts
layer(activation, [x], opts)
end
def activation(%Axon{output_shape: shape} = x, activation, opts)
when is_function(activation, 1) do
layer(activation, [x], [shape: shape] ++ opts)
end
## Activation
for {activation, name, a_or_an} <- @activation_layers do
@doc """
Adds #{a_or_an} #{name} activation layer to the network.
See `Axon.Activations.#{Atom.to_string(activation)}/1` for more details.
## Options
* `:name` - layer name.
"""
@doc type: :activation
def unquote(activation)(%Axon{} = x, opts \\ []) do
activation(x, unquote(activation), opts)
end
end
## Dropout
@dropout_layers [
{:dropout, "Dropout", "a"},
{:feature_alpha_dropout, "Feature alpha dropout", "a"},
{:spatial_dropout, "Spatial dropout", "a"},
{:alpha_dropout, "Alpha dropout", "an"}
]
for {dropout, name, a_or_an} <- @dropout_layers do
@doc """
Adds #{a_or_an} #{name} layer to the network.
See `Axon.Layers.#{Atom.to_string(dropout)}/2` for more details.
## Options
* `:name` - layer name.
* `:rate` - dropout rate. Defaults to `0.5`.
"""
@doc type: :dropout
def unquote(dropout)(%Axon{} = x, opts \\ []) do
dropout(x, unquote(dropout), opts)
end
end
defp dropout(%Axon{output_shape: parent_shape} = x, dropout, opts) do
opts = Keyword.validate!(opts, [:name, rate: 0.5])
layer(dropout, [x],
name: opts[:name],
rate: opts[:rate],
shape: parent_shape,
op_name: dropout
)
end
## Pooling
@pooling_layers [
{:max_pool, "Max pool", "a"},
{:avg_pool, "Average pool", "an"},
{:lp_pool, "Power average pool", "a"}
]
for {pool, name, a_or_an} <- @pooling_layers do
@doc """
Adds #{a_or_an} #{name} layer to the network.
See `Axon.Layers.#{Atom.to_string(pool)}/2` for more details.
## Options
* `:name` - layer name.
* `:kernel_size` - size of the kernel spatial dimensions. Defaults
to `1`.
* `:strides` - stride during convolution. Defaults to size of kernel.
* `:padding` - padding to the spatial dimensions of the input.
Defaults to `:valid`.
* `:dilations` - window dilations. Defaults to `1`.
* `:channels` - channels location. One of `:first` or `:last`.
Defaults to `:first`.
"""
@doc type: :pooling
def unquote(pool)(%Axon{} = x, opts \\ []) do
pool(x, unquote(pool), opts)
end
end
defp pool(%Axon{output_shape: parent_shape} = x, pool, opts) do
opts =
Keyword.validate!(opts, [
:name,
:strides,
kernel_size: 1,
padding: :valid,
channels: :first,
dilations: 1,
norm: 2
])
kernel_size = opts[:kernel_size]
strides = opts[:strides]
padding = opts[:padding]
channels = opts[:channels]
dilations = opts[:dilations]
inner_rank = Nx.rank(parent_shape) - 2
kernel_size = tuple_or_duplicate(:kernel_size, kernel_size, inner_rank)
strides = if strides, do: strides, else: Tuple.to_list(kernel_size)
strides = list_or_duplicate(:strides, strides, inner_rank)
dilations = list_or_duplicate(:dilations, dilations, inner_rank)
output_shape =
Axon.Shape.pool(parent_shape, kernel_size, strides, padding, dilations, channels)
name = opts[:name]
opts =
if pool == :lp_pool do
norm = opts[:norm]
[
name: name,
kernel_size: kernel_size,
strides: strides,
padding: padding,
channels: channels,
window_dilations: dilations,
norm: norm,
shape: output_shape,
op_name: pool
]
else
[
name: name,
kernel_size: kernel_size,
strides: strides,
padding: padding,
channels: channels,
window_dilations: dilations,
shape: output_shape,
op_name: pool
]
end
layer(pool, [x], opts)
end
## Adaptive Pooling
@adaptive_pooling_layers [
{:adaptive_avg_pool, "Adaptive average pool", "an"},
{:adaptive_max_pool, "Adaptive max pool", "an"},
{:adaptive_lp_pool, "Adaptive power average pool", "an"}
]
for {pool, name, a_or_an} <- @adaptive_pooling_layers do
@doc """
Adds #{a_or_an} #{name} layer to the network.
See `Axon.Layers.#{Atom.to_string(pool)}/2` for more details.
## Options
* `:name` - layer name.
* `:output_size` - layer output size.
* `:channels` - channel configuration. One of `:first` or `:last`.
Defaults to `:first`.
"""
@doc type: :pooling
def unquote(pool)(%Axon{} = x, opts \\ []) do
adaptative_pool(x, unquote(pool), opts)
end
end
defp adaptative_pool(%Axon{output_shape: parent_shape} = x, pool, opts) do
opts = Keyword.validate!(opts, [:name, :output_size, channels: :first, norm: 2])
channels = opts[:channels]
idx =
if channels == :first do
1
else
Nx.rank(parent_shape) - 1
end
output_size =
if size = opts[:output_size] do
size
else
parent_shape
|> Tuple.delete_at(0)
|> Tuple.delete_at(idx - 1)
end
inner_rank = Nx.rank(parent_shape) - 2
output_size = tuple_or_duplicate(:output_size, output_size, inner_rank)
output_shape = Axon.Shape.adaptive_pool(parent_shape, output_size, channels)
name = opts[:name]
opts =
if pool == :adaptive_lp_pool do
norm = opts[:norm]
[
name: name,
output_size: output_size,
norm: norm,
channels: channels,
shape: output_shape,
op_name: pool
]
else
[
name: name,
output_size: output_size,
channels: channels,
shape: output_shape,
op_name: pool
]
end
layer(pool, [x], opts)
end
## Global Pooling
@global_pooling_layers [
{:global_avg_pool, "Global average pool"},
{:global_max_pool, "Global max pool"},
{:global_lp_pool, "Global LP pool"}
]
for {pool, name} <- @global_pooling_layers do
@doc """
Adds a #{name} layer to the network.
See `Axon.Layers.#{Atom.to_string(pool)}/2` for more details.
Typically used to connect feature extractors such as those in convolutional
neural networks to fully-connected models by reducing inputs along spatial
dimensions to only feature and batch dimensions.
## Options
* `:name` - layer name.
* `:keep_axes` - option to keep reduced axes. If `true`, keeps reduced axes
with a dimension size of 1.
* `:channels` - channel configuration. One of `:first` or `:last`.
Defaults to `:first`.
"""
@doc type: :pooling
def unquote(pool)(%Axon{} = x, opts \\ []) do
global_pool(x, unquote(pool), opts)
end
end
defp global_pool(%Axon{output_shape: parent_shape} = x, pool, opts) do
opts = Keyword.validate!(opts, [:name, keep_axes: false, channels: :first, norm: 2])
keep_axes = opts[:keep_axes]
name = opts[:name]
channels = opts[:channels]
output_shape = Axon.Shape.global_pool(parent_shape, keep_axes, channels)
opts =
if pool == :global_lp_pool do
norm = opts[:norm]
[
name: name,
channels: channels,
keep_axes: keep_axes,
norm: norm,
shape: output_shape,
op_name: pool
]
else
[name: name, channels: channels, keep_axes: keep_axes, shape: output_shape, op_name: pool]
end
layer(pool, [x], opts)
end
## Normalization
@normalization_with_stats_layers [
{:batch_norm, "Batch normalization", "a"},
{:instance_norm, "Instance normalization", "an"}
]
for {norm, name, a_or_an} <- @normalization_with_stats_layers do
@doc """
Adds #{a_or_an} #{name} layer to the network.
See `Axon.Layers.#{Atom.to_string(norm)}/6` for more details.
## Options
* `:name` - layer name.
* `:gamma_initializer` - gamma parameter initializer. Defaults
to `:glorot_uniform`.
* `:beta_initializer` - beta parameter initializer. Defaults to
`:zeros`.
* `:channel_index` - input feature index used for calculating
mean and variance. Defaults to `1`.
* `:epsilon` - numerical stability term.
"""
@doc type: :normalization
def unquote(norm)(%Axon{} = x, opts \\ []) do
norm_with_stats(x, unquote(norm), opts)
end
end
defp norm_with_stats(%Axon{output_shape: shape} = x, norm, opts) do
opts =
Keyword.validate!(opts, [
:name,
gamma_initializer: :glorot_uniform,
beta_initializer: :zeros,
channel_index: 1,
epsilon: 1.0e-5,
momentum: 0.1
])
channel_index = opts[:channel_index]
gamma_shape = Axon.Shape.norm_param(shape, channel_index)
beta_shape = Axon.Shape.norm_param(shape, channel_index)
mean_shape = Axon.Shape.norm_param(shape, channel_index)
var_shape = Axon.Shape.norm_param(shape, channel_index)
gamma = param("gamma", gamma_shape, initializer: opts[:gamma_initializer])
beta = param("beta", beta_shape, initializer: opts[:beta_initializer])
mean = param("mean", mean_shape, initializer: :zeros)
var = param("var", var_shape, initializer: :ones)
layer(
norm,
[x, gamma, beta, mean, var],
name: opts[:name],
epsilon: opts[:epsilon],
channel_index: channel_index,
momentum: opts[:momentum],
shape: shape,
op_name: norm
)
end
@normalization_layers [
{:layer_norm, "Layer normalization", "a"}
]
for {norm, name, a_or_an} <- @normalization_layers do
@doc """
Adds #{a_or_an} #{name} layer to the network.
See `Axon.Layers.#{Atom.to_string(norm)}/4` for more details.
## Options
* `:name` - layer name.
* `:gamma_initializer` - gamma parameter initializer. Defaults
to `:glorot_uniform`.
* `:beta_initializer` - beta parameter initializer. Defaults to
`:zeros`.
* `:channel_index` - input feature index used for calculating
mean and variance. Defaults to `1`.
* `:epsilon` - numerical stability term.
"""
@doc type: :normalization
def unquote(norm)(%Axon{} = x, opts \\ []) do
norm(x, unquote(norm), opts)
end
end
defp norm(%Axon{output_shape: shape} = x, norm, opts) do
opts =
Keyword.validate!(opts, [
:name,
gamma_initializer: :glorot_uniform,
beta_initializer: :zeros,
channel_index: 1,
epsilon: 1.0e-5
])
channel_index = opts[:channel_index]
gamma_shape = Axon.Shape.norm_param(shape, channel_index)
beta_shape = Axon.Shape.norm_param(shape, channel_index)
gamma = param("gamma", gamma_shape, initializer: opts[:gamma_initializer])
beta = param("beta", beta_shape, initializer: opts[:beta_initializer])
layer(norm, [x, gamma, beta],
name: opts[:name],
epsilon: opts[:epsilon],
channel_index: channel_index,
shape: shape,
op_name: norm
)
end
@doc """
Adds a group normalization layer to the network.
See `Axon.Layers.group_norm/4` for more details.
## Options
* `:name` - layer name.
* `:gamma_initializer` - gamma parameter initializer. Defaults
to `:glorot_uniform`.
* `:beta_initializer` - beta parameter initializer. Defaults to
`:zeros`.
* `:channel_index` - input feature index used for calculating
mean and variance. Defaults to `1`.
* `:epsilon` - numerical stability term.
"""
@doc type: :normalization
def group_norm(%Axon{output_shape: shape} = x, group_size, opts \\ [])
when is_integer(group_size) and group_size >= 1 do
opts =
Keyword.validate!(opts, [
:name,
gamma_initializer: :glorot_uniform,
beta_initializer: :zeros,
channel_index: 1,
epsilon: 1.0e-5
])
channel_index = opts[:channel_index]
gamma_shape = Axon.Shape.norm_param(shape, channel_index)
beta_shape = Axon.Shape.norm_param(shape, channel_index)
gamma = param("gamma", gamma_shape, initializer: opts[:gamma_initializer])
beta = param("beta", beta_shape, initializer: opts[:beta_initializer])
layer(:group_norm, [x, gamma, beta],
name: opts[:name],
epsilon: opts[:epsilon],
channel_index: channel_index,
group_size: group_size,
shape: shape,
op_name: :group_norm
)
end
@doc """
Applies the given `Nx` expression to the input.
Nx layers are meant for quick applications of functions without
trainable parameters. For example, they are useful for applying
functions which apply accessors to containers:
model = Axon.container({foo, bar})
Axon.nx(model, &elem(&1, 0))
## Options
* `:name` - layer name.
"""
def nx(input, fun, opts \\ [])
@doc type: :special
def nx(%Axon{output_shape: input_shape} = x, fun, opts) when is_function(fun, 1) do
opts = Keyword.validate!(opts, [:name])
{name, opts} = Keyword.pop(opts, :name)
fun_with_params = fn x, _opts -> fun.(x) end
output_shape = infer_shape([input_shape], fun_with_params, opts)
layer(fun_with_params, [x], name: name, shape: output_shape, op_name: :nx)
end
@doc """
Adds a flatten layer to the network.
This layer will flatten all but the batch dimensions
of the input into a single layer. Typically called to flatten
the output of a convolution for use with a dense layer.
## Options
* `:name` - layer name.
* `:ignore_batch?` - whether to ignore batch dimension in
transpose operation. Defaults to `true`.
"""
@doc type: :shape
def flatten(%Axon{op: op, output_shape: shape} = x, opts \\ []) do
opts = Keyword.validate!(opts, [:name, ignore_batch?: op != :constant])
ignore_batch? = opts[:ignore_batch?]
output_shape = Axon.Shape.flatten(shape, ignore_batch?)
layer(:flatten, [x],
name: opts[:name],
ignore_batch?: ignore_batch?,
shape: output_shape,
op_name: :flatten
)
end
@doc """
Adds a reshape layer to the network.
This layer implements a special case of `Nx.reshape` which accounts
for possible batch dimensions in the input tensor. If the input contains
batch dimensions, the reshape operation is performed on all non-batch
dimensions of the input - preserving the original batch size.
If the input is an Axon constant, the reshape behavior matches that of
`Nx.reshape`.
## Options
* `:name` - layer name.
* `:ignore_batch?` - whether to ignore batch dimension in transpose
operation. Defaults to `true`.
"""
@doc type: :shape
def reshape(%Axon{op: op, output_shape: shape} = x, new_shape, opts \\ []) do
opts = Keyword.validate!(opts, [:name, ignore_batch?: op != :constant])
ignore_batch? = opts[:ignore_batch?]
output_shape = Axon.Shape.reshape(shape, new_shape, ignore_batch?)
layer(:reshape, [x],
name: opts[:name],
ignore_batch?: ignore_batch?,
shape: output_shape,
to: output_shape,
op_name: :reshape
)
end
@doc """
Adds a transpose layer to the network.
## Options
* `:name` - layer name.
* `:ignore_batch?` - whether to ignore batch dimension in transpose
operation. Defaults to true.
"""
@doc type: :shape
def transpose(%Axon{op: op, output_shape: shape} = x, permutation, opts \\ []) do
opts = Keyword.validate!(opts, [:name, ignore_batch?: op != :constant])
ignore_batch? = opts[:ignore_batch?]
output_shape = Axon.Shape.transpose(shape, permutation, ignore_batch?)
layer(:transpose, [x],
name: opts[:name],
axes: permutation,
ignore_batch?: ignore_batch?,
shape: output_shape,
op_name: :transpose
)
end
@doc """
Adds a pad layer to the network.
This layer will pad the spatial dimensions of the input.
Padding configuration is a list of tuples for each spatial
dimension.
## Options
* `:name` - layer name.
* `:channels` - channel configuration. One of `:first` or
`:last`. Defaults to `:first`.
"""
@doc type: :shape
def pad(%Axon{output_shape: shape} = x, config, value \\ 0.0, opts \\ [])
when is_list(config) and is_number(value) do
opts = Keyword.validate!(opts, [:name, channels: :first])
channels = opts[:channels]
output_shape = Axon.Shape.pad(shape, config)
layer(:pad, [x],
name: opts[:name],
padding_config: config,
value: value,
channels: channels,
shape: output_shape,
op_name: :pad
)
end
@doc """
Adds a resize layer to the network.
Resizing can be used for interpolation or upsampling input
values in a neural network. For example, you can use this
layer as an upsampling layer within a GAN.
Resize shape must be a tuple representing the resized spatial
dimensions of the input tensor.
Compiles to `Axon.Layers.resize/2`.
## Options
* `:name` - layer name.
* `:method` - resize method. Defaults to `:nearest`.
* `:channels` - channel configuration. One of `:first` or
`:last`. Defaults to `:first`.
"""
@doc type: :shape
def resize(%Axon{output_shape: shape} = x, resize_shape, opts \\ []) do
opts = Keyword.validate!(opts, [:name, method: :nearest, channels: :first])
channels = opts[:channels]
output_shape = Axon.Shape.resize(shape, resize_shape, channels)
layer(:resize, [x],
name: opts[:name],
method: opts[:method],
channels: channels,
shape: output_shape,
to: resize_shape,
op_name: :resize
)
end
@doc """
Adds a concatenate layer to the network.
This layer will concatenate inputs along the last
dimension unless specified otherwise.
## Options
* `:name` - layer name.
* `:axis` - concatenate axis. Defaults to `-1`.
"""
@doc type: :combinator
def concatenate(%Axon{output_shape: x_shape} = x, %Axon{output_shape: y_shape} = y, opts)
when is_list(opts) do
opts = Keyword.validate!(opts, [:name, axis: -1])
axis = opts[:axis]
output_shape = Axon.Shape.concatenate([x_shape, y_shape], axis)
layer(:concatenate, [container({x, y})],
name: opts[:name],
axis: axis,
shape: output_shape,
op_name: :concatenate
)
end
@doc type: :combinator
def concatenate([%Axon{} | _] = inputs, opts)
when is_list(inputs) and is_list(opts) do
opts = Keyword.validate!(opts, [:name, axis: -1])
axis = opts[:axis]
input_shapes = inputs |> Enum.map(fn %Axon{output_shape: shape} -> shape end)
output_shape = Axon.Shape.concatenate(input_shapes, axis)
layer(:concatenate, [container(List.to_tuple(inputs))],
name: opts[:name],
axis: axis,
shape: output_shape,
op_name: :concatenate
)
end
@doc false
def concatenate(%Axon{} = x, %Axon{} = y), do: concatenate(x, y, [])
@doc false
def concatenate(inputs) when is_list(inputs), do: concatenate(inputs, [])
@element_wise_layers [:add, :subtract, :multiply]
for op <- @element_wise_layers do
@doc """
Adds a #{op} layer to the network.
This layer performs an element-wise #{Atom.to_string(op)} operation
on input layers. All input layers must be capable of being
broadcast together.
If one shape has a static batch size, all other shapes must have a
static batch size as well.
## Options
* `:name` - layer name.
"""
@doc type: :combinator
def unquote(op)(%Axon{output_shape: lhs_shape} = x, %Axon{output_shape: rhs_shape} = y, opts) do
opts = Keyword.validate!(opts, [:name])
output_shape = Axon.Shape.element_wise([lhs_shape, rhs_shape])
layer(unquote(op), [container({x, y})],
name: opts[:name],
shape: output_shape,
op_name: unquote(op)
)
end
@doc """
Adds a #{op} layer to the network.
This layer performs an element-wise #{Atom.to_string(op)} operation
on all input layers. All input layers must be capable of being
broadcast together.
## Options
* `:name` - layer name.
"""
@doc type: :combinator
def unquote(op)(inputs, opts) when is_list(inputs) and is_list(opts) do
opts = Keyword.validate!(opts, [:name])
shapes =
Enum.map(inputs, fn
%Axon{output_shape: shape} -> shape
invalid -> raise ArgumentError, "invalid input #{inspect(invalid)}"
end)
output_shape = Axon.Shape.element_wise(shapes)
layer(unquote(op), [container(List.to_tuple(inputs))],
name: opts[:name],
shape: output_shape,
op_name: unquote(op)
)
end
@doc false
def unquote(op)(%Axon{} = x, %Axon{} = y) do
unquote(op)(x, y, [])
end
@doc false
def unquote(op)([%Axon{} | _] = inputs), do: unquote(op)(inputs, [])
end
@doc """
Adds a conditional layer which conditionally executes
`true_graph` or `false_graph` based on the condition `cond_fn`
at runtime.
`cond_fn` is an arity-1 function executed on the output of the
parent graph. It must return a boolean scalar tensor (e.g. 1 or 0).
The shapes of `true_graph` and `false_graph` must be equal.
"""
@doc type: :combinator
def cond(
%Axon{} = parent,
cond_fn,
%Axon{output_shape: out_shape} = true_graph,
%Axon{output_shape: out_shape} = false_graph,
opts \\ []
)
when is_function(cond_fn, 1) do
opts = Keyword.validate!(opts, [:name])
layer(:cond, [parent, true_graph, false_graph],
name: opts[:name],
cond: cond_fn,
shape: out_shape,
op_name: :cond
)
end
@doc """
Splits input graph into a container of `n` input graphs
along the given axis.
## Options
* `:name` - layer name.
* `:axis` - concatenate axis. Defaults to `-1`.
"""
@doc type: :combinator
def split(parent, splits, opts \\ [])
def split(%Axon{} = parent, splits, opts) when is_list(splits) do
opts = Keyword.validate!(opts, [:name, axis: -1])
axis = opts[:axis]
{_, split_layers} =
for {split, i} <- Enum.with_index(splits), reduce: {0, []} do
{num_split, split_layers} ->
name =
case opts[:name] do
names when is_list(names) ->
Enum.at(names, i)
name ->
name
end
layer =
layer(
fn x, _ -> Nx.slice_along_axis(x, num_split, split, axis: axis) end,
[parent],
name: name,
op_name: :split
)
{num_split + split, [layer | split_layers]}
end
split_layers |> Enum.reverse() |> List.to_tuple()
end
def split(%Axon{output_shape: shape} = parent, n, opts) when is_integer(n) do
opts = Keyword.validate!(opts, [:name, axis: -1])
axis = opts[:axis]
{slice_size, split_shape} = Axon.Shape.split(shape, n, axis)
splits =
for i <- 0..(n - 1) do
name =
case opts[:name] do
names when is_list(names) ->
Enum.at(names, i)
name ->
name
end
layer(
fn x, _ -> Nx.slice_along_axis(x, i * slice_size, slice_size, axis: axis) end,
[parent],
name: name,
shape: split_shape,
op_name: :split
)
end
List.to_tuple(splits)
end
@doc """
See `lstm/3`.
"""
@doc type: :recurrent
def lstm(%Axon{} = x, units) when is_integer(units) and units > 0 do
lstm(x, units, [])
end
@doc """
Adds a long short-term memory (LSTM) layer to the network
with a random initial hidden state.
See `lstm/4` for more details.
## Additional options
* `:recurrent_initializer` - initializer for hidden state.
Defaults to `:glorot_uniform`.
"""
@doc type: :recurrent
def lstm(%Axon{output_shape: shape} = x, units, opts)
when is_integer(units) and units > 0 and is_list(opts) do
{recurrent_initializer, opts} = Keyword.pop(opts, :recurrent_initializer, :glorot_uniform)
c = rnn_state(x, shape, units, :lstm, opts[:name], "c", recurrent_initializer)
h = rnn_state(x, shape, units, :lstm, opts[:name], "h", recurrent_initializer)
lstm(x, {c, h}, units, opts)
end
def lstm(%Axon{} = x, {%Axon{}, %Axon{}} = hidden_state, units)
when is_integer(units) and units > 0 do
lstm(x, hidden_state, units, [])
end
@doc """
Adds a long short-term memory (LSTM) layer to the network
with the given initial hidden state.
LSTMs apply `Axon.Recurrent.lstm_cell/7` over an entire input
sequence and return:
{{new_cell, new_hidden}, output_sequence}
You can use the output state as the hidden state of another
LSTM layer.
## Options
* `:name` - layer name.
* `:activation` - recurrent activation. Defaults to `:tanh`.
* `:gate` - recurrent gate function. Defaults to `:sigmoid`.
* `:unroll` - `:dynamic` (loop preserving) or `:static` (compiled)
unrolling of RNN.
* `:kernel_initializer` - initializer for kernel weights. Defaults
to `:glorot_uniform`.
* `:bias_initializer` - initializer for bias weights. Defaults to
`:zeros`.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`.
"""
@doc type: :recurrent
def lstm(
%Axon{output_shape: shape} = x,
{%Axon{output_shape: h_shape}, %Axon{output_shape: h_shape}} = hidden_state,
units,
opts \\ []
)
when is_integer(units) and units > 0 and is_list(opts) do
opts =
Keyword.validate!(opts, [
:name,
activation: :tanh,
gate: :sigmoid,
unroll: :dynamic,
use_bias: true,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros
])
activation = opts[:activation]
gate = opts[:gate]
unroll = opts[:unroll]
output_shape = Axon.Shape.rnn(shape, units, :lstm)
input_kernel_shape = Axon.Shape.rnn_input_kernel(shape, units, :lstm)
hidden_kernel_shape = Axon.Shape.rnn_hidden_kernel(shape, units, :lstm)
bias_shape = Axon.Shape.rnn_bias(shape, units, :lstm)
kernel_initializer = opts[:kernel_initializer]
# Parameters
input_kernel =
param("input_kernel", {:tuple, List.duplicate(input_kernel_shape, 4)},
initializer: kernel_initializer
)
hidden_kernel =
param("hidden_kernel", {:tuple, List.duplicate(hidden_kernel_shape, 4)},
initializer: kernel_initializer
)
hidden_state_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"lstm_#{op_counts[:lstm]}_hidden_state"
end
name when is_binary(name) ->
"#{name}_hidden_state"
end
hidden_state = Axon.container(hidden_state, name: hidden_state_name)
{inputs, op} =
if opts[:use_bias] do
bias_initializer = opts[:bias_initializer]
bias =
param("bias", {:tuple, List.duplicate(bias_shape, 4)}, initializer: bias_initializer)
{[x, hidden_state, input_kernel, hidden_kernel, bias], :lstm}
else
{[x, hidden_state, input_kernel, hidden_kernel], &Axon.Layers.lstm(&1, &2, &3, &4, 0, &5)}
end
output =
layer(
op,
inputs,
name: opts[:name],
activation: activation,
gate: gate,
unroll: unroll,
shape: {{h_shape, h_shape}, output_shape},
op_name: :lstm
)
new_c_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"lstm_#{op_counts[:lstm]}_c_hidden_state"
end
name when is_binary(name) ->
"#{name}_c_hidden_state"
end
new_h_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"lstm_#{op_counts[:lstm]}_h_hidden_state"
end
name when is_binary(name) ->
"#{name}_h_hidden_state"
end
output_sequence_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"lstm_#{op_counts[:lstm]}_output_sequence"
end
name when is_binary(name) ->
"#{name}_output_sequence"
end
new_c =
layer(fn x, _ -> elem(elem(x, 0), 0) end, [output],
name: new_c_name,
shape: h_shape,
op_name: :elem
)
new_h =
layer(fn x, _ -> elem(elem(x, 0), 1) end, [output],
name: new_h_name,
shape: h_shape,
op_name: :elem
)
output_sequence =
layer(fn x, _ -> elem(x, 1) end, [output],
name: output_sequence_name,
shape: output_shape,
op_name: :elem
)
{{new_c, new_h}, output_sequence}
end
@doc """
See `gru/3`.
"""
@doc type: :recurrent
def gru(%Axon{} = x, units) do
gru(x, units, [])
end
@doc """
Adds a gated recurrent unit (GRU) layer to the network with
a random initial hidden state.
See `gru/4` for more details.
## Additional options
* `:recurrent_initializer` - initializer for hidden state.
Defaults to `:glorot_uniform`.
"""
@doc type: :recurrent
def gru(%Axon{output_shape: shape} = x, units, opts)
when is_integer(units) and units > 0
when is_list(opts) do
{recurrent_initializer, opts} = Keyword.pop(opts, :recurrent_initializer, :glorot_uniform)
h = rnn_state(x, shape, units, :gru, opts[:name], "h", recurrent_initializer)
gru(x, {h}, units, opts)
end
def gru(%Axon{} = x, {%Axon{}} = hidden_state, units) when is_integer(units) and units > 0 do
gru(x, hidden_state, units, [])
end
@doc """
Adds a gated recurrent unit (GRU) layer to the network with
the given initial hidden state.
GRUs apply `Axon.Recurrent.gru_cell/7` over an entire input
sequence and return:
{{new_hidden}, output_sequence}
You can use the output state as the hidden state of another
GRU layer.
## Options
* `:name` - layer name.
* `:activation` - recurrent activation. Defaults to `:tanh`.
* `:gate` - recurrent gate function. Defaults to `:sigmoid`.
* `:unroll` - `:dynamic` (loop preserving) or `:static` (compiled)
unrolling of RNN.
* `:kernel_initializer` - initializer for kernel weights. Defaults
to `:glorot_uniform`.
* `:bias_initializer` - initializer for bias weights. Defaults to
`:zeros`.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`.
"""
@doc type: :recurrent
def gru(
%Axon{output_shape: shape} = x,
{%Axon{output_shape: h_shape}} = hidden_state,
units,
opts
)
when is_integer(units) and units > 0 and is_list(opts) do
opts =
Keyword.validate!(opts, [
:name,
activation: :tanh,
gate: :sigmoid,
unroll: :dynamic,
use_bias: true,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros
])
activation = opts[:activation]
gate = opts[:gate]
unroll = opts[:unroll]
output_shape = Axon.Shape.rnn(shape, units, :gru)
input_kernel_shape = Axon.Shape.rnn_input_kernel(shape, units, :gru)
hidden_kernel_shape = Axon.Shape.rnn_hidden_kernel(shape, units, :gru)
bias_shape = Axon.Shape.rnn_bias(shape, units, :gru)
kernel_initializer = opts[:kernel_initializer]
input_kernel =
param("input_kernel", {:tuple, List.duplicate(input_kernel_shape, 3)},
initializer: kernel_initializer
)
hidden_kernel =
param("hidden_kernel", {:tuple, List.duplicate(hidden_kernel_shape, 3)},
initializer: kernel_initializer
)
hidden_state_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"gru_#{op_counts[:gru]}_hidden_state"
end
name when is_binary(name) ->
"#{name}_hidden_state"
end
hidden_state = Axon.container(hidden_state, name: hidden_state_name)
inputs =
if opts[:use_bias] do
bias_initializer = opts[:bias_initializer]
bias =
param("bias", {:tuple, List.duplicate(bias_shape, 4)}, initializer: bias_initializer)
[x, hidden_state, input_kernel, hidden_kernel, bias]
else
[x, hidden_state, input_kernel, hidden_kernel]
end
output =
layer(
:gru,
inputs,
name: opts[:name],
activation: activation,
gate: gate,
unroll: unroll,
shape: {{h_shape}, output_shape},
op_name: :gru
)
new_h_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"gru_#{op_counts[:gru]}_hidden_state"
end
name when is_binary(name) ->
"#{name}_hidden_state"
end
output_sequence_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"gru_#{op_counts[:gru]}_output_sequence"
end
name when is_binary(name) ->
"#{name}_output_sequence"
end
new_h =
layer(fn x, _ -> elem(elem(x, 0), 0) end, [output],
name: new_h_name,
shape: h_shape,
op_name: :elem
)
output_sequence =
layer(fn x, _ -> elem(x, 1) end, [output],
name: output_sequence_name,
shape: output_shape,
op_name: :elem
)
{{new_h}, output_sequence}
end
@doc """
See `conv_lstm/3`.
"""
@doc type: :recurrent
def conv_lstm(%Axon{} = x, units) when is_integer(units) and units > 0 do
conv_lstm(x, units, [])
end
@doc """
Adds a convolutional long short-term memory (LSTM) layer to the network
with a random initial hidden state.
See `conv_lstm/4` for more details.
## Additional options
* `:recurrent_initializer` - initializer for hidden state. Defaults
to `:glorot_uniform`.
"""
@doc type: :recurrent
def conv_lstm(%Axon{output_shape: shape} = x, units, opts)
when is_integer(units) and units > 0 and is_list(opts) do
{recurrent_initializer, opts} = Keyword.pop(opts, :recurrent_initializer, :glorot_uniform)
c = rnn_state(x, shape, units, :conv_lstm, opts[:name], "c", recurrent_initializer)
h = rnn_state(x, shape, units, :conv_lstm, opts[:name], "h", recurrent_initializer)
conv_lstm(x, {c, h}, units, opts)
end
def conv_lstm(%Axon{} = x, {%Axon{}, %Axon{}} = hidden_state, units)
when is_integer(units) and units > 0 do
conv_lstm(x, hidden_state, units, [])
end
@doc """
Adds a convolutional long short-term memory (LSTM) layer to the network
with the given initial hidden state..
ConvLSTMs apply `Axon.Recurrent.conv_lstm_cell/5` over an entire input
sequence and return:
{{new_cell, new_hidden}, output_sequence}
You can use the output state as the hidden state of another
ConvLSTM layer.
## Options
* `:name` - layer name.
* `:padding` - convolutional padding. Defaults to `:same`.
* `:kernel_size` - convolutional kernel size. Defaults to `1`.
* `:strides` - convolutional strides. Defaults to `1`.
* `:unroll` - `:dynamic` (loop preserving) or `:static` (compiled)
unrolling of RNN.
* `:kernel_initializer` - initializer for kernel weights. Defaults
to `:glorot_uniform`.
* `:bias_initializer` - initializer for bias weights. Defaults to
`:zeros`.
* `:use_bias` - whether the layer should add bias to the output.
Defaults to `true`.
"""
@doc type: :recurrent
def conv_lstm(
%Axon{output_shape: shape} = x,
{%Axon{output_shape: h_shape}, %Axon{output_shape: h_shape}} = hidden_state,
units,
opts
)
when is_integer(units) and units > 0 and is_list(opts) do
opts =
Keyword.validate!(opts, [
:name,
padding: :same,
kernel_size: 1,
strides: 1,
unroll: :dynamic,
kernel_initializer: :glorot_uniform,
bias_initializer: :zeros,
use_bias: true
])
padding = opts[:padding]
kernel_size = opts[:kernel_size]
strides = opts[:strides]
unroll = opts[:unroll]
inner_rank = Nx.rank(shape) - 3
sequence_length = elem(shape, 1)
kernel_size = tuple_or_duplicate(:kernel_size, kernel_size, inner_rank)
strides = list_or_duplicate(:strides, strides, inner_rank)
input_dilation = List.duplicate(1, inner_rank)
kernel_dilation = List.duplicate(1, inner_rank)
conv_shape = Tuple.delete_at(shape, 1)
conv_hidden_state_shape = Tuple.delete_at(h_shape, 1)
hidden_kernel_shape =
Axon.Shape.conv_kernel(conv_hidden_state_shape, 4 * units, kernel_size, :first)
input_kernel_shape = Axon.Shape.conv_kernel(conv_shape, 4 * units, kernel_size, :first)
bias_shape = Axon.Shape.conv_bias(conv_shape, 4 * units, kernel_size, :first)
output_kernel_shape =
Axon.Shape.conv_kernel(conv_hidden_state_shape, units, kernel_size, :first)
output_shape =
conv_hidden_state_shape
|> Axon.Shape.conv(
output_kernel_shape,
strides,
padding,
input_dilation,
kernel_dilation,
:first,
1
)
|> Tuple.insert_at(1, sequence_length)
kernel_initializer = opts[:kernel_initializer]
wi = param("input_kernel", {:tuple, [input_kernel_shape]}, initializer: kernel_initializer)
wh = param("hidden_kernel", {:tuple, [hidden_kernel_shape]}, initializer: kernel_initializer)
hidden_state_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"conv_lstm_#{op_counts[:conv_lstm]}_hidden_state"
end
name when is_binary(name) ->
"#{name}_hidden_state"
end
hidden_state = Axon.container(hidden_state, name: hidden_state_name)
{inputs, op} =
if opts[:use_bias] do
bias_initializer = opts[:bias_initializer]
b = param("bias", {:tuple, [bias_shape]}, initializer: bias_initializer)
{[x, hidden_state, wi, wh, b], :conv_lstm}
else
{[x, hidden_state, wi, wh], &Axon.Layers.conv_lstm(&1, &2, &3, &4, {0}, &5)}
end
output =
layer(
op,
inputs,
name: opts[:name],
conv_opts: [
strides: strides,
padding: padding
],
unroll: unroll,
shape: output_shape,
op_name: :conv_lstm
)
new_c_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"conv_lstm_#{op_counts[:lstm]}_c_hidden_state"
end
name when is_binary(name) ->
"#{name}_c_hidden_state"
end
new_h_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"conv_lstm_#{op_counts[:lstm]}_h_hidden_state"
end
name when is_binary(name) ->
"#{name}_h_hidden_state"
end
output_sequence_name =
case opts[:name] do
nil ->
fn _, op_counts ->
"conv_lstm_#{op_counts[:lstm]}_output_sequence"
end
name when is_binary(name) ->
"#{name}_output_sequence"
end
new_c =
layer(fn x, _ -> elem(elem(x, 0), 0) end, [output],
name: new_c_name,
shape: h_shape,
op_name: :elem
)
new_h =
layer(fn x, _ -> elem(elem(x, 0), 1) end, [output],
name: new_h_name,
shape: h_shape,
op_name: :elem
)
output_sequence =
layer(fn x, _ -> elem(x, 1) end, [output],
name: output_sequence_name,
shape: output_shape,
op_name: :elem
)
{{new_c, new_h}, output_sequence}
end
defp rnn_state(x, shape, units, rnn_type, parent_name, state_name, initializer) do
initializer = initializer || :glorot_uniform
name =
case parent_name do
nil ->
fn _, op_counts ->
"lstm_#{op_counts[rnn_type]}_#{state_name}_hidden_state"
end
parent_name when is_binary(parent_name) ->
"#{parent_name}_#{state_name}_hidden_state"
end
shape = Axon.Shape.rnn_hidden_state(shape, units, rnn_type)
fun = fn inputs, _opts ->
shape = put_elem(shape, 0, elem(Nx.shape(inputs), 0))
case initializer do
fun when is_function(fun) ->
fun.(shape)
fun when is_atom(fun) ->
fun = apply(Axon.Initializers, fun, [])
fun.(shape, {:f, 32})
end
end
layer(fun, [x], name: name, op_name: :recurrent_state)
end
@doc """
Adds an embedding layer to the network.
An embedding layer initializes a kernel of shape `{vocab_size, embedding_size}`
which acts as a lookup table for sequences of discrete tokens (e.g. sentences).
Embeddings are typically used to obtain a dense representation of a sparse input
space.
## Options
* `:name` - layer name.
* `:kernel_initializer` - initializer for `kernel` weights. Defaults
to `:uniform`.
"""
@doc type: :linear
def embedding(%Axon{output_shape: shape} = x, vocab_size, embedding_size, opts \\ []) do
opts = Keyword.validate!(opts, [:name, kernel_initializer: :uniform])
kernel_shape = Axon.Shape.embedding_kernel(shape, vocab_size, embedding_size)
output_shape = Axon.Shape.embedding(shape, vocab_size, embedding_size)
kernel = param("kernel", kernel_shape, initializer: opts[:kernel_initializer])
layer(:embedding, [x, kernel], name: opts[:name], shape: output_shape, op_name: :embedding)
end
@doc """
Adds a bias layer to the network.
A bias layer simply adds a trainable bias to an input.
## Options
* `:name` - layer name.
* `:bias_initializer` - initializer for `bias` weights. Defaults
to `:zeros`.
"""
@doc type: :linear
def bias(%Axon{output_shape: shape} = x, opts \\ []) do
opts = Keyword.validate!(opts, [:name, bias_initializer: :zeros])
units = elem(shape, tuple_size(shape) - 1)
bias_shape = Axon.Shape.dense_bias(shape, units)
bias = param("bias", bias_shape, initializer: opts[:bias_initializer])
layer(:bias, [x, bias], name: opts[:name], shape: shape, op_name: :bias)
end
@doc """
Freezes parameters returned from `fun` in the given
model. `fun` takes the model's parameter list and returns
the list of parameters it wishes to freeze. `fun` defaults
to the identity function, freezing all of the parameters in
`model`.
Freezing parameters is useful when performing transfer learning
to leverage features learned from another problem in a new problem.
For example, it's common to combine the convolutional base from
larger models trained on ImageNet with fresh fully-connected classifiers.
The combined model is then trained on fresh data, with the convolutional
base frozen so as not to lose information. You can see this example in code
here:
cnn_base = get_pretrained_cnn_base()
model =
cnn_base
|> Axon.freeze()
|> Axon.flatten()
|> Axon.dense(1024, activation: :relu)
|> Axon.dropout()
|> Axon.dense(1000, activation: :softmax)
model
|> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adam(0.005))
|> Axon.Loop.run(data, epochs: 10)
When compiled, frozen parameters are wrapped in `Nx.Defn.Kernel.stop_grad/1`,
which zeros out the gradient with respect to the frozen parameter. Gradients
of frozen parameters will return `0.0`, meaning they won't be changed during
the update process.
"""
def freeze(%Axon{} = model, fun \\ & &1) when is_function(fun, 1) do
parameters =
tree_reduce(model, MapSet.new(), fn %Axon{parameters: params}, acc ->
Enum.reduce(params, acc, fn param, acc ->
MapSet.put(acc, param)
end)
end)
parameters_to_freeze = fun.(Enum.to_list(parameters))
tree_map(model, fn %Axon{parameters: params} = axon ->
frozen_params =
Enum.map(params, fn %{name: param_name} = v ->
if Enum.any?(parameters_to_freeze, fn %{name: name} -> name == param_name end) do
%{v | frozen: true}
else
v
end
end)
%{axon | parameters: frozen_params}
end)
end
@doc """
Attaches a hook to the given Axon model.
Hooks compile down to `Nx.Defn.Kernel.hook/3` and provide the same
functionality for adding side-effecting operations to a compiled
model. For example, you can use hooks to inspect intermediate activations,
send data to an external service, and more.
Hooks can be configured to be invoked on the following events:
* `:initialize` - on model initialization.
* `:pre_forward` - before layer forward pass is invoked.
* `:forward` - after layer forward pass is invoked.
* `:backward` - after layer backward pass is invoked.
To invoke a hook on every single event, you may pass `:all` to `on:`.
Axon.input({nil, 1}, "input") |> Axon.attach_hook(&IO.inspect/1, on: :all)
The default event is `:forward`, assuming you want a hook invoked
on the layers forward pass.
You may configure hooks to run in one of only training or inference
mode using the `:mode` option. The default mode is `:both` to be invoked
during both train and inference mode.
Axon.input({nil, 1}, "input") |> Axon.attach_hook(&IO.inspect/1, on: :forward, mode: :train)
You can also attach multiple hooks to a single layer. Hooks are invoked in
the order in which they are declared. If order is important, you should attach
hooks in the order you want them to be executed:
Axon.input({nil, 1}, "input")
# I will be executed first
|> Axon.attach_hook(&IO.inspect/1)
# I will be executed second
|> Axon.attach_hook(fn _ -> IO.write("HERE") end)
Hooks are executed at their point of attachment. You must insert hooks at each point
you want a hook to execute during model execution.
Axon.input({nil, 1}, "input")
|> Axon.attach_hook(&IO.inspect/1)
|> Axon.relu()
|> Axon.attach_hook(&IO.inspect/1)
"""
def attach_hook(%Axon{hooks: hooks} = axon, fun, opts \\ []) do
opts = Keyword.validate!(opts, on: :forward, mode: :both)
on_event = opts[:on]
mode = opts[:mode]
%{axon | hooks: [{on_event, mode, fun} | hooks]}
end
## Traversal
@doc """
Traverses a model tree applying `fun` to each layer.
"""
def tree_map(%Axon{op: :container, parent: [container]} = axon, fun) do
x = deep_new(container, fun)
%{fun.(axon) | parent: [x]}
end
def tree_map(%Axon{parent: x} = axon, fun) when is_list(x) do
x = Enum.map(x, &tree_map(&1, fun))
%{fun.(axon) | parent: x}
end
@doc """
Traverses a model applying `fun` with an accumulator.
"""
def tree_reduce(%Axon{op: :container, parent: [container]} = axon, acc, fun) do
deep_reduce(container, fun.(axon, acc), fun)
end
def tree_reduce(%Axon{parent: x} = axon, acc, fun) when is_list(x) do
Enum.reduce(x, fun.(axon, acc), &tree_reduce(&1, &2, fun))
end
# TODO: Should not be duplicated
def deep_reduce(map, acc, fun) do
Nx.Container.reduce(map, acc, &recur_deep_reduce(&1, &2, fun))
end
defp recur_deep_reduce(value, acc, fun) do
case value do
%Axon{} = val ->
fun.(val, acc)
%Nx.Tensor{} = val ->
fun.(val, acc)
{:leaf, val} ->
fun.(val, acc)
val ->
deep_reduce(val, acc, fun)
end
end
## Utilities
@doc """
Returns the model's signature as a tuple of `{input_shape, output_shape}`.
## Examples
iex> model = Axon.input({nil, 32}, "input") |> Axon.dense(10)
iex> {inp, out} = Axon.get_model_signature(model)
iex> inp
{nil, 32}
iex> out
{nil, 10}
iex> inp1 = Axon.input({nil, 32}, "input_0")
iex> inp2 = Axon.input({nil, 32}, "input_1")
iex> model = Axon.concatenate(inp1, inp2)
iex> {{inp1_shape, inp2_shape}, out} = Axon.get_model_signature(model)
iex> inp1_shape
{nil, 32}
iex> inp2_shape
{nil, 32}
iex> out
{nil, 64}
"""
def get_model_signature(%Axon{output_shape: output_shape} = axon) do
# TODO: Refactor for tuples and use `tree_*` when they support
# tuple inputs
input_shapes =
tree_reduce(axon, [], fn
%Axon{op: :input, output_shape: shape}, acc -> [shape | acc]
_, acc -> acc
end)
case input_shapes do
[input_shape] ->
{input_shape, output_shape}
shapes ->
{List.to_tuple(Enum.reverse(shapes)), output_shape}
end
end
@doc """
Compiles the given model to `{init_fn, predict_fn}`.
Once compiled, a model can be passed as argument to `Nx.Defn`.
"""
@doc type: :compilation
def compile(model, opts \\ []) when is_list(opts) do
{Axon.Compiler.compile_init(model, opts), Axon.Compiler.compile_predict(model, opts)}
end
@doc """
Compiles and runs the given models initialization function
with the given compiler options.
You may optionally specify initial parameters for some layers or
namespaces by passing a partial parameter map:
Axon.init(model, %{"dense_0" => dense_params})
The parameter map will be merged with the initialized model
parameters.
"""
@doc type: :execution
def init(model, params \\ %{}, opts \\ []) when is_list(opts) do
Axon.Compiler.compile_init(model, opts).(params)
end
@doc """
Compiles and runs the given Axon model with `params` on
`input` with the given compiler options.
"""
@doc type: :execution
def predict(%Axon{} = model, params, input, opts \\ []) when is_list(opts) do
Axon.Compiler.compile_predict(model, opts).(params, input)
end
## Inspection
defimpl Inspect do
import Inspect.Algebra
import Axon.Shared
alias Axon.Parameter
def inspect(axon, _opts) do
title = "Model"
header = ["Layer", "Shape", "Policy", "Parameters", "Parameters Memory"]
model_info = %{num_params: 0, total_param_byte_size: 0, inputs: []}
{_, _, cache, _, model_info} = axon_to_rows(axon, %{}, %{}, model_info)
rows =
cache
|> Enum.sort()
|> Enum.unzip()
|> elem(1)
|> Enum.map(&elem(&1, 0))
rows
|> TableRex.Table.new(header, title)
|> TableRex.Table.render!(
header_separator_symbol: "=",
title_separator_symbol: "=",
vertical_style: :off
)
|> then(&(&1 <> "Total Parameters: #{model_info.num_params}\n"))
|> then(&(&1 <> "Total Parameters Memory: #{model_info.total_param_byte_size} bytes\n"))
|> then(&(&1 <> "Inputs: #{inspect(Map.new(model_info.inputs))}\n"))
|> string()
end
defp axon_to_rows(%{id: id, op_name: op_name} = graph, cache, op_counts, model_info) do
case cache do
%{^id => {row, name}} ->
{row, name, cache, op_counts, model_info}
%{} ->
{row, name, cache, op_counts, model_info} =
do_axon_to_rows(graph, cache, op_counts, model_info)
cache = Map.put(cache, id, {row, name})
op_counts = Map.update(op_counts, op_name, 1, fn x -> x + 1 end)
{row, name, cache, op_counts, model_info}
end
end
defp do_axon_to_rows(
%Axon{
op: :container,
parent: [parents],
name: name_fn,
output_shape: shape,
policy: policy
},
cache,
op_counts,
model_info
) do
{input_names, {cache, op_counts, model_info}} =
deep_map_reduce(parents, {cache, op_counts, model_info}, fn
graph, {cache, op_counts, model_info} ->
{_, name, cache, op_counts, model_info} =
axon_to_rows(graph, cache, op_counts, model_info)
{name, {cache, op_counts, model_info}}
end)
op_string = "container"
name = name_fn.(:container, op_counts)
row = [
"#{name} ( #{op_string} #{inspect(input_names)} )",
"#{inspect(shape)}",
"#{inspect(policy)}",
0,
"0 bytes"
]
{row, name, cache, op_counts, model_info}
end
defp do_axon_to_rows(
%Axon{
op: :namespace,
parent: parents,
name: name_fn,
output_shape: shape,
policy: policy
},
cache,
op_counts,
model_info
) do
init_model_info = %{num_params: 0, total_param_byte_size: 0, inputs: []}
{_input_names, {_cache, op_counts, namespace_model_info}} =
Enum.map_reduce(parents, {%{}, op_counts, init_model_info}, fn
graph, {cache, op_counts, model_info} ->
{_, name, cache, op_counts, model_info} =
axon_to_rows(graph, cache, op_counts, model_info)
{name, {cache, op_counts, model_info}}
end)
name = name_fn.(:namespace, op_counts)
num_params = namespace_model_info.num_params
param_byte_size = namespace_model_info.total_param_byte_size
inputs = namespace_model_info.inputs
model_info =
model_info
|> Map.update(:num_params, 0, fn x -> x + num_params end)
|> Map.update(:total_param_byte_size, 0, fn x -> x + param_byte_size end)
|> Map.update(:inputs, [], fn x -> x ++ inputs end)
row = [
"#{name} ( #{inputs |> Map.new() |> Map.keys()} )",
"#{inspect(shape)}",
"#{inspect(policy)}",
"#{num_params}",
"#{param_byte_size} bytes"
]
{row, name, cache, op_counts, model_info}
end
defp do_axon_to_rows(
%Axon{
parent: parents,
parameters: params,
name: name_fn,
output_shape: shape,
policy: %{params: {_, bitsize}} = policy,
op_name: op_name
},
cache,
op_counts,
model_info
) do
{input_names, {cache, op_counts, model_info}} =
Enum.map_reduce(parents, {cache, op_counts, model_info}, fn
graph, {cache, op_counts, model_info} ->
{_, name, cache, op_counts, model_info} =
axon_to_rows(graph, cache, op_counts, model_info)
{name, {cache, op_counts, model_info}}
end)
num_params =
Enum.reduce(params, 0, fn
%Parameter{shape: {:tuple, shapes}}, acc ->
Enum.reduce(shapes, acc, &(Nx.size(&1) + &2))
%Parameter{shape: shape}, acc ->
acc + Nx.size(shape)
end)
param_byte_size = num_params * div(bitsize, 8)
op_inspect = Atom.to_string(op_name)
inputs =
case input_names do
[] ->
""
[_ | _] = input_names ->
"#{inspect(input_names)}"
end
name = name_fn.(op_name, op_counts)
row = [
"#{name} ( #{op_inspect}#{inputs} )",
"#{inspect(shape)}",
"#{inspect(policy)}",
"#{num_params}",
"#{param_byte_size} bytes"
]
model_info =
model_info
|> Map.update(:num_params, 0, &(&1 + num_params))
|> Map.update(:total_param_byte_size, 0, &(&1 + param_byte_size))
|> Map.update(:inputs, [], fn inputs ->
if op_name == :input, do: [{name, shape} | inputs], else: inputs
end)
{row, name, cache, op_counts, model_info}
end
end
## Serialization
@doc """
Serializes a model and its parameters for persisting
models to disk or elsewhere.
Model and parameters are serialized as a tuple, where the
model is converted to a recursive map to ensure compatibility
with future Axon versions and the parameters are serialized
using `Nx.serialize/2`. There is some additional metadata included
such as current serialization version for compatibility.
Serialization `opts` are forwarded to `Nx.serialize/2` and
`:erlang.term_to_binary/2` for controlling compression options.
## Examples
iex> model = Axon.input({nil, 2}, "input") |> Axon.dense(1, kernel_initializer: :zeros, activation: :relu)
iex> params = Axon.init(model)
iex> serialized = Axon.serialize(model, params)
iex> {saved_model, saved_params} = Axon.deserialize(serialized)
iex> Axon.predict(saved_model, saved_params, Nx.tensor([[1.0, 1.0]]))
#Nx.Tensor<
f32[1][1]
[
[0.0]
]
>
"""
def serialize(%Axon{} = model, params, opts \\ []) do
model_meta = axon_to_map(model)
params = Nx.serialize(params, opts)
:erlang.term_to_binary({@file_version, model_meta, params}, opts)
end
defp axon_to_map(%Axon{op: :container, parent: [parents]} = model) do
parents = deep_new(parents, &axon_to_map/1)
axon_map = Map.from_struct(model) |> Map.put(:axon, :axon)
%{axon_map | parent: List.wrap(parents)}
end
defp axon_to_map(%Axon{parent: parents} = model) do
parents = Enum.map(parents, &axon_to_map/1)
axon_map = Map.from_struct(model) |> Map.put(:axon, :axon)
%{axon_map | parent: parents}
end
@doc """
Deserializes serialized model and parameters into a `{model, params}`
tuple.
It is the opposite of `Axon.serialize/3`.
## Examples
iex> model = Axon.input({nil, 2}, "input") |> Axon.dense(1, kernel_initializer: :zeros, activation: :relu)
iex> params = Axon.init(model)
iex> serialized = Axon.serialize(model, params)
iex> {saved_model, saved_params} = Axon.deserialize(serialized)
iex> Axon.predict(saved_model, saved_params, Nx.tensor([[1.0, 1.0]]))
#Nx.Tensor<
f32[1][1]
[
[0.0]
]
>
"""
def deserialize(serialized, opts \\ []) do
{1, model_meta, serialized_params} = :erlang.binary_to_term(serialized, [:safe | opts])
model = map_to_axon(model_meta)
params = Nx.deserialize(serialized_params, opts)
{model, params}
end
defp map_to_axon(%{op: :container, parent: [parents]} = model) do
parents = deep_new(parents, &map_to_axon/1)
model = Map.drop(model, [:axon])
model = %{model | parent: List.wrap(parents)}
struct(__MODULE__, model)
end
defp map_to_axon(%{axon: :axon, parent: parents} = model) do
parents = Enum.map(parents, &map_to_axon/1)
model = Map.drop(model, [:axon])
model = %{model | parent: parents}
struct(__MODULE__, model)
end
## Helpers
@valid_initializers [:zeros, :ones, :uniform, :normal, :identity] ++
[:lecun_uniform, :lecun_normal, :he_uniform, :he_normal] ++
[:glorot_uniform, :glorot_normal, :variance_scaling]
defp validate_initializer!(initializer)
when is_atom(initializer) and initializer in @valid_initializers do
:ok
end
defp validate_initializer!(initializer) when is_function(initializer, 2) do
:ok
end
defp validate_initializer!(initializer) do
raise ArgumentError,
"initializer must be one of #{inspect(@valid_initializers)}," <>
" or an arity-2 function accepting initializer shape and type" <>
" got #{inspect(initializer)}"
end
defp tuple_or_duplicate(key, tuple_or_integer, rank) do
cond do
is_tuple(tuple_or_integer) ->
if tuple_size(tuple_or_integer) != rank do
raise ArgumentError,
"expected #{inspect(key)} to be a #{rank}-element tuple, " <>
"got: #{inspect(tuple_or_integer)}"
end
tuple_or_integer
is_integer(tuple_or_integer) ->
Tuple.duplicate(tuple_or_integer, rank)
true ->
raise ArgumentError,
"expected #{inspect(key)} to be an integer or a tuple, " <>
"got: #{inspect(tuple_or_integer)}"
end
end
defp list_or_duplicate(key, list_or_integer, rank) do
cond do
is_list(list_or_integer) ->
if length(list_or_integer) != rank do
raise ArgumentError,
"expected #{inspect(key)} to be a #{rank}-element list, " <>
"got: #{inspect(list_or_integer)}"
end
list_or_integer
is_integer(list_or_integer) ->
List.duplicate(list_or_integer, rank)
true ->
raise ArgumentError,
"expected #{inspect(key)} to be an integer or a list, " <>
"got: #{inspect(list_or_integer)}"
end
end
# Names are generated lazily at inspect, initialization, and compile
# time, so for name we return a function which takes `op` and `op_count`
# and returns a unique name for the given model.
defp unique_identifiers(type, nil) do
id = System.unique_integer([:positive, :monotonic])
name = fn op, op_counts ->
count = op_counts[op] || 0
Atom.to_string(type) <> "_#{count}"
end
{id, name}
end
defp unique_identifiers(_type, name_fn) when is_function(name_fn, 2) do
id = System.unique_integer([:positive, :monotonic])
{id, name_fn}
end
defp unique_identifiers(_type, name) when is_binary(name) do
{System.unique_integer([:positive, :monotonic]), fn _, _ -> name end}
end
defp unique_identifiers(_, name) do
raise ArgumentError,
"expected layer name to be a binary, a function or nil, " <>
"got: #{inspect(name)}"
end
end
