Powered by AppSignal & Oban Pro
Would you like to see your link here? Contact us
Notesclub
Sign in with GitHub
created by hec & contributors
terms privacy

Elixir Axon tutorial

langs/elixir/axon/testing.livemd

avatar
José Luis
avatar
avatar
Share to X
Share to Bluesky
More notebooks

Elixir Axon tutorial

Mix.install([
  {:nx, "~> 0.2.0"},
  {:exla, "~> 0.2.0"},
  {:axon, "~> 0.1.0-dev", github: "elixir-nx/axon", branch: "main"}
])

Aim

We will create and train a neural network to detect oddity on one byte numbers

Create a neural network

Let’s start creating a neural network for our aim

require Axon

model =
  Axon.input({nil, 8})
  |> Axon.dense(20, activation: :relu)
  |> Axon.dense(20, activation: :relu)
  # |> Axon.dropout(rate: 0.5)
  |> Axon.dense(2, activation: :softmax)

Easy-peasy

As you can see, its a network with 8 neurons as input

  Axon.input({nil, 8})

Later, we put two small layers of neurons (just 20 neurons per layer, it looks an easy problem and probably is not necessary more)

  |> Axon.dense(20, activation: :relu)
  |> Axon.dense(20, activation: :relu)

And for finish, just two neurons as output

  |> Axon.dense(2, activation: :softmax)

And… that’s all we already have our Neuron network

With a little initialization, we already can ask to our network

Let’s do it!

state_init = model |> Axon.init(compiler: EXLA)

Now we have a neural network initialized and we can ask to it

Let’s start asking for number 0

Axon.predict(
  model,
  state_init,
  %{
    "input_0" => Nx.tensor([[0.0, 0.0, 0, 0, 0, 0, 0.0, 0.0]])
  },
  compiler: EXLA
)

The values of the tensor should say as the probability of be odd or even

> [0.5, 0.5]

As you can see, the system has no clue about the answer

And with number 1, similar result

It’s quite logical, at the moment we only have a neural network with a specific topology but we didn’t add to the network any information about odd and even numbers

How to?

Trainning

Prepare data to train

First we need data to train

It consist on input values, with the right response

Let’s start with something simple

train_data = [
  {Nx.tensor([[0.0, 0, 0, 0, 0, 0, 0, 1.0]]), Nx.tensor([0.0, 1.0])},
  {Nx.tensor([[0.0, 0, 0, 0, 0, 0, 0.0, 0.0]]), Nx.tensor([1.0, 0.0])}
]

And update the neural network with these training data

trained0 =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005))
  |> Axon.Loop.run(train_data, epochs: 100, compiler: EXLA)

Now we have a neural network trainned with a little information (but we repeated a lot).

Lets ask to our NN again

Axon.predict(
  model,
  trained0,
  %{
    "input_0" => Nx.tensor([[0.0, 0.0, 0, 0, 0, 0, 0.0, 0.0]])
  },
  compiler: EXLA
)
Axon.predict(
  model,
  trained0,
  %{
    "input_0" => Nx.tensor([[0.0, 0.0, 0, 0, 0, 0, 0.0, 1.0]])
  },
  compiler: EXLA
)

That looks quite different!!!

But, what if we ask a not trainned number?

Axon.predict(
  model,
  trained0,
  %{
    "input_0" => Nx.tensor([[0, 0, 0, 1, 0, 0, 0, 0]])
  },
  compiler: EXLA
)

Not very good answer

We have to train with several numbers, let’s do it!

To do that, we want a big (infinity will be enough ;-) list of random numbers with the right response.

random = Enum.random(0..255)
even = if rem(random, 2) == 0, do: true, else: false
{random, even}

Let’s put this in an infinite list

trainning_data =
  Stream.unfold(0, fn _ ->
    random = Enum.random(0..255)
    even = if rem(random, 2) == 0, do: true, else: false
    {{random, even}, 0}
  end)

Show some data to verify

trainning_data |> Enum.take(10)

Looks great, but… remember how we introduced our data in the network

train_data = [
  {Nx.tensor([[0.0, 0, 0, 0, 0, 0, 0, 1.0]]), Nx.tensor([0.0, 1.0])},
  {Nx.tensor([[0.0, 0, 0, 0, 0, 0, 0.0, 0.0]]), Nx.tensor([1.0, 0.0])}
]

Then we have to map to this format

First, a function to convert a number, to a tensor with the neurons to be activated

number2input_tensor = fn num ->
  :io_lib.format("~8.2.0B", [num])
  |> Enum.map(fn d -> if d == ?0, do: 0.0, else: 1.0 end)
  |> Nx.tensor()
  |> Nx.reshape({1, 8})
end

number2input_tensor.(123)

And now, a function to convert true or false to the tensor result

even2output_tensor = fn is_even ->
  if(is_even, do: Nx.tensor([1.0, 0.0]), else: Nx.tensor([0.0, 1.0]))
end

even2output_tensor.(false)

Remember we started…

trainning_data |> Enum.take(10)
trainning_data =
  trainning_data
  |> Stream.map(fn {n, is_even} ->
    {number2input_tensor.(n), even2output_tensor.(is_even)}
  end)

trainning_data |> Enum.take(10)

Trainning day!

Now we have a infinite list of data to train!

Just train it again with some datas

trained1 =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005))
  |> Axon.Loop.run(trainning_data |> Stream.take(1000), epochs: 1, compiler: EXLA)

Section

Section

Section

Section

require Axon

model =
  Axon.input({nil, 8})
  |> Axon.dense(20, activation: :relu)
  |> Axon.dense(20, activation: :relu)
  # |> Axon.dropout(rate: 0.5)
  |> Axon.dense(2, activation: :softmax)

# |> Axon.dense(2, activation: :sigmoid)

# state_init = model |> Axon.init(compiler: EXLA)

train_data = [
  {Nx.tensor([[0.0, 0, 0, 0, 0, 0, 0, 1.0]]), Nx.tensor([0.0, 1.0])},
  {Nx.tensor([[0.0, 0, 0, 0, 0, 0, 0.0, 0.0]]), Nx.tensor([1.0, 0.0])}
]

trained =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005))
  |> Axon.Loop.run(train_data, epochs: 100, compiler: EXLA)

Axon.predict(model, trained, %{
  "input_0" => Nx.tensor([[0.0, 0.0, 0, 0, 0, 0, 0.0, 1.0]])
})

Testing

require Axon

model =
  Axon.input({nil, 8})
  |> Axon.dense(128)
  |> Axon.dense(2, activation: :softmax)
state_init = model |> Axon.init(compiler: EXLA)
train_data = [
  {Nx.tensor([[0.0, 0, 0, 0, 0, 0, 0, 1.0]]), Nx.tensor([0.0, 1.0])},
  {Nx.tensor([[0.0, 0, 0, 0, 0, 0, 1.0, 0.0]]), Nx.tensor([1.0, 0.0])}
]

model_state =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005))
  |> Axon.Loop.run(train_data, epochs: 10, compiler: EXLA)
Axon.predict(model, model_state, %{
  "input_0" => Nx.tensor([[0.0, 0.0, 0, 0, 0, 0, 1.0, 0.0]])
})
:io_lib.format("~8.2.0B", [1])
|> Enum.map(fn d -> if d == ?0, do: 0.0, else: 1.0 end)
|> Nx.tensor()
random_num_oddity = fn ->
  random = Enum.random(0..255)
  odd = if rem(random, 2) == 0, do: true, else: false

  bin_tensor =
    :io_lib.format("~8.2.0B", [random])
    |> Enum.map(fn d -> if d == ?0, do: 0.0, else: 1.0 end)
    |> Nx.tensor()
    |> Nx.reshape({1, 8})

  {bin_tensor, odd}
end

{test_1, test_2} = random_num_oddity.()
random_num_oddity_list = fn ->
  Stream.unfold(0, fn _ ->
    {random_num_oddity.(), 0}
  end)
end

random_num_oddity_list.() |> Enum.take(10)
data_trainning =
  Stream.unfold(0, fn _ ->
    {random_num_oddity.(), 0}
  end)
  |> Stream.map(fn {tn, oddity} ->
    {tn, if(oddity, do: Nx.tensor([0.0, 1.0]), else: Nx.tensor([1.0, 0.0]))}
  end)

# |> Enum.take(10)
require Axon

model =
  Axon.input({nil, 8})
  |> Axon.dense(20, activation: :relu)
  |> Axon.dense(20, activation: :relu)
  # |> Axon.dropout(rate: 0.5)
  |> Axon.dense(2, activation: :softmax)

# |> Axon.dense(2, activation: :sigmoid)

trained =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005))
  |> Axon.Loop.run(data_trainning |> Enum.take(1_000), epochs: 1, compiler: EXLA)

Axon.predict(
  model,
  trained,
  %{
    "input_0" => Nx.tensor([[0.0, 0.0, 0, 0, 0, 0, 0.0, 1.0]])
  },
  compiler: EXLA
)
Nx.tensor([[0.0, 0, 0, 0, 0, 0, 0, 1.0]])

Other notebooks:

Back