Neural Network In Axon

Mix.install([
  {:axon, "~> 0.6.1"},
  {:nx, "~> 0.7.2"},
  {:exla, "~> 0.7.2"},
  {:kino, "~> 0.12.3"},
  {:scidata, "~> 0.1.11"}
])

MNIST Dataset

THE MNIST DATABASE of handwritten digits

base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"
train_image_file = "train-images-idx3-ubyte.gz"
train_label_file = "train-labels-idx1-ubyte.gz"
test_image_file = "t10k-images-idx3-ubyte.gz"
test_label_file = "t10k-labels-idx1-ubyte.gz"

train_image_url = base_url <> train_image_file
train_label_url = base_url <> train_label_file
test_image_url = base_url <> test_image_file
test_label_url = base_url <> test_label_file

Train Image Data

{:ok, {_status, _headers, train_image_file}} = :httpc.request(train_image_url)

TRAINING SET IMAGE FILE (train-images-idx3-ubyte)

<<_magic_number::32, n_images::32, n_rows::32, n_columns::32, train_image_file::binary>> = :zlib.gunzip(train_image_file)

train_image_tensor =
  train_image_file
  |> Nx.from_binary({:u, 8})
  |> Nx.reshape({n_images, n_rows, n_columns})
  |> Nx.divide(255)

train_image_tensor |> Nx.to_heatmap()

Train label Data

{:ok, {_status, _headers, train_label_file }} = :httpc.request(train_label_url)

TRAINING SET LABEL FILE (train-labels-idx1-ubyte)

The labels values are 0 to 9.

<<_::32, n_labels::32, train_label_file::binary>> = :zlib.gunzip(train_label_file)

train_label_tensor =
  train_label_file
  |> Nx.from_binary({:u, 8})
   |> Nx.new_axis(-1)
  |> Nx.equal(Nx.tensor(Enum.to_list(0..9)))

Neural Network

model =
  Axon.input("input", shape: {32, 28 , 28})
  |> Axon.flatten()
  |> Axon.dense(128, activation: :sigmoid)
  |> Axon.dense(10, activation: :softmax)

# Axon.Display.as_graph(model, Nx.template({nil, 1, 28, 28}, {:u, 8}))

train_image_batch = Nx.to_batched(train_image_tensor, 32)

train_label_batch = Nx.to_batched(train_label_tensor, 32)

Stream.zip(train_image_batch, train_label_batch) |> Enum.to_list()

params =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, :adam)
  |> Axon.Loop.run(Stream.zip(train_image_batch, train_label_batch), %{},
    epochs: 10,
    compiler: EXLA
  )

first_batch = Enum.at(train_image_batch, 0)

output = Axon.predict(model, params, first_batch)

Nx.argmax(output, axis: 1)