Hacker's Handbook

Neural Networks in Elixir

A beginners guide

Posted: 2023-06-19

Today was an exciting day as I embarked on a journey to explore the world of neural networks using Elixir. My goal is to leverage the power of the BEAM virtual machine and the functional programming paradigm to build a neural network model and for fun implement the layers in the network as a network of communicating processes.

Neural networks are a subset of machine learning algorithms modeled after the human brain. They are designed to recognize patterns and interpret sensory data through a kind of machine perception, labeling or clustering raw input. These algorithms can be used to recognize complex patterns, make predictions, or make decisions in a diverse range of applications including image and speech recognition, medical diagnosis, statistical arbitrage, learning to rank, and even in games. The networks themselves are composed of interconnected nodes or 'neurons', which are organized into layers. Each layer processes the input it receives and passes on a transformed version of the input to the next layer. This structure allows neural networks to learn from data and improve over time, making them a powerful tool in the field of artificial intelligence.

As it turns out the great developers in the Elixir community has already done almost all the work. I have been meaning to try this out for a long time but have not had the time to do it. But today I had a Sunday afternoon free and could choose between playing Diablo IV or doing som coding for fun.

I hope to revisit this topic in a future blog post, where I plan to enhance Axon's functionality by using processes and process communication instead of Elixir streams for the loops I don't expect it to be more efficient but I think it would be a fun and cool project to do.

Todays blog will be more of a beginners guide on how to do a basic set up and how to test that it works.

I am assuming lots of things, for example that you are running Ubuntu linux and have some development experience. This is more a guide for me, than for you as a reader, sorry ;).

Setting Up the Environment

The first step was to set up the development environment. I used asdf, a version manager that can manage multiple language runtime versions on a per-project basis.

Note that these steps are specific to Ubuntu and that I am running Ubuntu in WSL2 under Windows 11, your mileage may vary.

I cloned the asdf repository and added it to my shell profile:

git clone https://github.com/asdf-vm/asdf.git ~/.asdf --branch v0.12.0
echo -e '\n. "$HOME/.asdf/asdf.sh"' >> ~/.bashrc
echo -e '\n. "$HOME/.asdf/completions/asdf.bash"' >> ~/.bashrc
. ~/.bashrc

Next, I added Erlang and Elixir plugins to asdf and installed the versions I needed:

asdf plugin add erlang
asdf plugin add elixir
asdf install erlang 26.0.1
asdf install elixir main-otp-26

Creating a New Elixir Project

With the environment set up, I created a new Elixir project using the mix new command:

mix new test_nx
cd test_nx

Adding Dependencies

I added the necessary dependencies to the mix.exs file. I used Nx, a multi-dimensional tensors library for Elixir, and Torchx, a library that provides bindings for the LibTorch tensor library.

Nx is the Elixir tensor computation library that allows you to perform numerical computations, which are essential for building neural networks.

A tensor is a mathematical object that generalizes the concepts of scalars, vectors, and matrices to higher dimensions. It's essentially an array of numbers, arranged on a grid, with a variable number of axes. Tensors provide a natural and compact way of representing data and transformations of data in machine learning, physics, and many other fields. In the context of machine learning, tensors are particularly important because they efficiently represent and manipulate multi-dimensional data structures, such as images, which can be represented as 3D tensors (height, width, color channels), or text, which can be represented as 2D tensors (sequence length, word embeddings). The ability to work with tensors allows us to handle a wide range of complex data types, making them a fundamental building block in these fields.

Torchx is an Elixir binding to the popular Torch machine learning library which I in the future will use to bring in some of my old models from PyTorch to run on the BEAM instead.

defp deps do
  [
    {:nx, "~> 0.5"},
    {:torchx, "~> 0.5"}
  ]
end

Then, I fetched the dependencies:

mix deps.get

Exploring Nx

I started an IEx session with my Mix dependencies:

iex -S mix

I experimented with Nx by creating a tensor and performing some operations on it copied directly from the Nx documentation

iex> t = Nx.tensor([[1, 2], [3, 4]])
iex> Nx.divide(Nx.exp(t), Nx.sum(Nx.exp(t)))
#Nx.Tensor<
  f32[2][2]
  [
    [0.032058604061603546, 0.08714432269334793],
    [0.23688282072544098, 0.6439142227172852]
  ]
>

It worked!

Building a Neural Network with Axon

Axon is a deep learning library for the Elixir programming language. It provides a high-level, flexible API for defining neural network models in Elixir. Axon leverages the Nx library for tensor computations, which allows for efficient numerical computation that is essential in machine learning. The reason for using Axon in this context is its seamless integration with Elixir and the BEAM virtual machine, providing a functional approach to defining and training neural networks. It also supports automatic differentiation and GPU acceleration, which are crucial for training complex models efficiently.

I added Axon and Exla (and stb_image for a later step, and another blog) to my dependencies:

 defp deps do
    [
      # {:dep_from_hexpm, "~> 0.3.0"},
      # {:dep_from_git, git: "https://github.com/elixir-lang/my_dep.git", tag: "0.1.0"}
      {:nx, "~> 0.5"},
      {:torchx, "~> 0.5"},
      {:axon, "~> 0.5"},
      {:exla, "~> 0.5"},
      {:stb_image, "~> 0.5.2"}
    ]

Downloaded the dependencies and started Elixir:

mix deps.get
iex -S mix

To test this I used an example from the Axon examples.

This model is a simple XOR function that takes two binary inputs and outputs a single binary value. The model is trained using a batch of random binary inputs and their corresponding XOR results.

The XOR function, a fundamental operation in computing and digital logic, stands for "exclusive or", and it's a binary operation that takes two inputs. If exactly one of the inputs is true (or 1 in binary terms), it returns true (or 1). In all other cases (both inputs are false or both are true), it returns false (or 0).

In the context of this guide, the XOR function is used as a simple problem for the neural network to learn. It's a non-linear problem, meaning it can't be solved using simple linear methods, which makes it a good test for a neural network. The goal is to train the network to understand the XOR logic, i.e., given two binary inputs, it should output the correct XOR result. This serves as a basic proof of concept that the neural network is functioning correctly and can learn to approximate functions.

defmodule XOR do

require Axon

  defp build_model(input_shape1, input_shape2) do
    inp1 = Axon.input("x1", shape: input_shape1)
    inp2 = Axon.input("x2", shape: input_shape2)

    inp1
    |> Axon.concatenate(inp2)
    |> Axon.dense(8, activation: :tanh)
    |> Axon.dense(1, activation: :sigmoid)
  end

  defp batch do
    x1 = Nx.tensor(for _ <- 1..32, do: [Enum.random(0..1)])
    x2 = Nx.tensor(for _ <- 1..32, do: [Enum.random(0..1)])
    y = Nx.logical_xor(x1, x2)
    { %{"x1" => x1, "x2" => x2}, y }
  end

  defp train_model(model, data, epochs) do
    model
    |> Axon.Loop.trainer(:binary_cross_entropy, :sgd)
    |> Axon.Loop.run(data, %{}, epochs: epochs, iterations: 1000, compiler: EXLA)
  end

  def run do
    model = build_model({nil, 1}, {nil, 1})
    data = Stream.repeatedly(&batch/0)

    model_state = train_model(model, data, 10)

    IO.inspect(
      Axon.predict(model, model_state, %{"x1" => Nx.tensor([[0]]), "x2" => Nx.tensor([[1]])})
    )
  end

end

XOR.run()

Wrapping Up

After running the model, I got the following output:

18:39:39.689 [debug] Forwarding options: [compiler: EXLA] to JIT compiler

18:39:39.729 [info] TfrtCpuClient created.
Epoch: 0, Batch: 950, loss: 0.6493790
Epoch: 1, Batch: 950, loss: 0.5824046
Epoch: 2, Batch: 950, loss: 0.5003822
Epoch: 3, Batch: 950, loss: 0.4245099
Epoch: 4, Batch: 950, loss: 0.3641715
Epoch: 5, Batch: 950, loss: 0.3175022
Epoch: 6, Batch: 950, loss: 0.2810990
Epoch: 7, Batch: 950, loss: 0.2520537
Epoch: 8, Batch: 950, loss: 0.2284686
Epoch: 9, Batch: 950, loss: 0.2089621
#Nx.Tensor<
  f32[1][1]
  EXLA.Backend<host:0, 0.3871477933.15073302.157805>
  [
    [0.9698225259780884]
  ]
>

The model successfully learned the XOR function and was able to predict the output for the inputs [0] and [1].

It's so nice not to have to use Python for writing and running a model!

Today was a productive day, full of learning and exploration. I'm excited to continue my journey into the world of neural networks with Elixir.

Next, I will see if I can get CUDA to work so I can use my NVIDIA 3090 card for something useful...

... .But that's a story for another day.