Would you like to see your link here? Contact us

Notesclub

created by hec & contributors

terms privacy

Machine Learning in Elixir (Ch 6+)

ml_ch6.livemd

@aos

gpt-from-scratch-nx-axon

More notebooks

Machine Learning in Elixir (Ch 6+)

Mix.install(
  [
    {:axon, "~> 0.5"},
    {:axon_onnx, github: "mortont/axon_onnx", override: true},
    {:nx, "~> 0.5"},
    {:polaris, "~> 0.1"},
    {:explorer, "~> 0.5"},
    {:kino, "~> 0.8"},
    {:scholar, "~> 0.3.1"},
    {:exla, "~> 0.5"},
    {:benchee, github: "bencheeorg/benchee", override: true},
    {:table_rex, "~> 3.1.1"},
    {:scidata, "~> 0.1"},
    {:stb_image, "~> 0.6"},
    {:vega_lite, "~> 0.1"},
    {:kino_vega_lite, "~> 0.1"}
  ],
  config: [
    nx: [
      default_backend: {EXLA.Backend, client: :cuda},
      default_defn_options: [compiler: EXLA, client: :cuda]
    ],
    exla: [
      clients: [
        cuda: [platform: :cuda, memory_fraction: 0.25]
      ]
    ]
  ],
  system_env: [
    XLA_TARGET: "cuda120",
  ]
)

Setup v2

require Explorer.DataFrame, as: DF
alias VegaLite, as: Vl

Chapter 6

Neural network

defmodule NeuralNetwork do
  import Nx.Defn

  defn hidden(input, weight, bias) do
    input
    |> dense(weight, bias)
    |> activation()
  end

  defn output(input, weight, bias) do
    input
    |> dense(weight, bias)
    |> activation()
  end

  defn predict(input, w1, b1, w2, b2) do
    input
    |> hidden(w1, b1)
    |> output(w2, b2)
  end

  defn dense(input, weight, bias) do
    input
    |> Nx.dot(weight)
    |> Nx.add(bias)
  end

  defn activation(input) do
    Nx.sigmoid(input)
  end
end

key = Nx.Random.key(42)
{w1, new_key} = Nx.Random.uniform(key)
{b1, new_key} = Nx.Random.uniform(new_key)
{w2, new_key} = Nx.Random.uniform(new_key)
{b2, new_key} = Nx.Random.uniform(new_key)

{input, _new_key} = Nx.Random.uniform(new_key, shape: {})

input
|> NeuralNetwork.predict(w1, b2, w2, b2)

Axon

{images, labels} = Scidata.MNIST.download()

{image_data, image_type, image_shape} = images
{label_data, label_type, label_shape} = labels

images =
  image_data
  |> Nx.from_binary(image_type)
  |> Nx.divide(255)
  |> Nx.reshape({60000, :auto})

labels =
  label_data
  |> Nx.from_binary(label_type)
  |> Nx.reshape(label_shape)
  |> Nx.new_axis(-1)
  |> Nx.equal(Nx.iota({1, 10}))

train_range = 0..49_999//1
test_range = 50_000..-1//1

train_images = images[train_range]
train_labels = labels[train_range]

test_images = images[test_range]
test_labels = labels[test_range]

batch_size = 64

train_data =
  train_images
  |> Nx.to_batched(batch_size)
  |> Stream.zip(Nx.to_batched(train_labels, batch_size))

test_data =
  test_images
  |> Nx.to_batched(batch_size)
  |> Stream.zip(Nx.to_batched(test_labels, batch_size))

Building the model with Axon

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

template = Nx.template({1, 784}, :f32)
Axon.Display.as_graph(model, template)

Axon.Display.as_table(model, template)
|> IO.puts()

IO.inspect(model, structs: false)

Training the model

trained_model_state =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, :sgd)
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.run(train_data, %{}, epochs: 10, compiler: EXLA)

Evaluating the model

model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(test_data, trained_model_state, compiler: EXLA)

Executing models with Axon

{test_batch, _} = Enum.at(test_data, 0)
test_image = test_batch[0]

test_image
|> Nx.reshape({28, 28})
|> Nx.to_heatmap()

{_, predict_fn} = Axon.build(model, compiler: EXLA)

probabilities =
  test_image
  |> Nx.new_axis(0)
  |> then(&amp;predict_fn.(trained_model_state, &amp;1))

probabilities |> Nx.argmax()

Chapter 7

Creating a pipeline

defmodule CatsAndDogs do
  def pipeline(paths, batch_size, target_height, target_width) do
    paths
    |> Enum.shuffle()
    |> Task.async_stream(&amp;parse_image/1)
    |> Stream.filter(fn
      {:ok, {%StbImage{}, _}} -> true
      _ -> false
    end)
    |> Stream.map(&amp;to_tensors(&amp;1, target_height, target_width))
    |> Stream.chunk_every(batch_size, batch_size, :discard)
    |> Stream.map(fn chunks ->
      {img_chunk, label_chunk} = Enum.unzip(chunks)
      {Nx.stack(img_chunk), Nx.stack(label_chunk)}
    end)
  end

  def pipeline_with_aug(paths, batch_size, target_height, target_width) do
    paths
    |> Enum.shuffle()
    |> Task.async_stream(&amp;parse_image/1)
    |> Stream.filter(fn
      {:ok, {%StbImage{}, _}} -> true
      _ -> false
    end)
    |> Stream.map(&amp;to_tensors(&amp;1, target_height, target_width))
    |> Stream.map(&amp;random_flip(&amp;1, :height))
    |> Stream.map(&amp;random_flip(&amp;1, :width))
    |> Stream.chunk_every(batch_size, batch_size, :discard)
    |> Stream.map(fn chunks ->
      {img_chunk, label_chunk} = Enum.unzip(chunks)
      {Nx.stack(img_chunk), Nx.stack(label_chunk)}
    end)
  end

  defp parse_image(path) do
    base = Path.basename(path)
    label = if String.contains?(base, "cat"), do: 0, else: 1

    case StbImage.read_file(path) do
      {:ok, img} -> {img, label}
      _error -> :error
    end
  end

  defp to_tensors({:ok, {img, label}}, target_height, target_width) do
    img_tensor =
      img
      |> StbImage.resize(target_height, target_width)
      |> StbImage.to_nx()
      |> Nx.divide(255)

    label_tensor = Nx.tensor([label])

    {img_tensor, label_tensor}
  end

  defp random_flip({image, label}, axis) do
    if :rand.uniform() < 0.5 do
      {Nx.reverse(image, axes: [axis]), label}
    else
      {image, label}
    end
  end
end

batch_size = 128
target_height = 96
target_width = 96

{test_paths, train_paths} =
  Path.wildcard("train/cats_dogs/*.jpg")
  |> Enum.shuffle()
  |> Enum.split(1000)

{test_paths, val_paths} = test_paths |> Enum.split(750)

train_pipeline = CatsAndDogs.pipeline_with_aug(
  train_paths, batch_size, target_height, target_width
)
val_pipeline = CatsAndDogs.pipeline(
  val_paths, batch_size, target_height, target_width
)
test_pipeline = CatsAndDogs.pipeline(
  test_paths, batch_size, target_height, target_width
)

Enum.take(train_pipeline, 1)

Training the MLP

mlp_model =
  Axon.input("images", shape: {nil, target_height, target_width, 3})
  |> Axon.flatten()
  |> Axon.dense(256, activation: :relu)
  |> Axon.dense(128, activation: :relu)
  |> Axon.dense(1, activation: :sigmoid)

mlp_trained_model_state =
  mlp_model
  |> Axon.Loop.trainer(:binary_cross_entropy, :adam)
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.run(train_pipeline, %{}, epochs: 5, compiler: EXLA)

mlp_model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(test_pipeline, mlp_trained_model_state, compiler: EXLA)

Convolutional Networks

path = "train/cats_dogs/dog.5.jpg"
img =
  path
  |> StbImage.read_file!()
  |> StbImage.to_nx()
  |> Nx.transpose(axes: [:channels, :height, :width])
  |> Nx.new_axis(0)

kernel = Nx.tensor([
  [-1, 0, 1],
  [-1, 0, 1],
  [-1, 0, 1]
])
kernel = kernel |> Nx.reshape({1, 1, 3, 3}) |> Nx.broadcast({3, 3, 3, 3})

img
|> Nx.conv(kernel)
|> Nx.as_type({:u, 8})
|> Nx.squeeze(axes: [0])
|> Nx.transpose(axes: [:height, :width, :channels])
|> Kino.Image.new()

Implementing CNNs

cnn_model =
  Axon.input("images", shape: {nil, 96, 96, 3})
  |> Axon.conv(32, kernel_size: {3, 3}, activation: :relu, padding: :same)
  |> Axon.batch_norm()
  |> Axon.max_pool(kernel_size: {2, 2}, strides: [2, 2])
  |> Axon.conv(64, kernel_size: {3, 3}, activation: :relu, padding: :same)
  |> Axon.batch_norm()
  |> Axon.max_pool(kernel_size: {2, 2}, strides: [2, 2])
  |> Axon.conv(128, kernel_size: {3, 3}, activation: :relu, padding: :same)
  |> Axon.max_pool(kernel_size: {2, 2}, strides: [2, 2])
  |> Axon.flatten()
  |> Axon.dense(128, activation: :relu)
  |> Axon.dropout(rate: 0.5)
  |> Axon.dense(1, activation: :sigmoid)

template = Nx.template({1, 96, 96, 3}, :f32)
Axon.Display.as_graph(cnn_model, template)

cnn_trained_model_state =
  cnn_model
  |> Axon.Loop.trainer(:binary_cross_entropy, Polaris.Optimizers.adam(learning_rate: 1.0e-3))
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.validate(cnn_model, val_pipeline)
  |> Axon.Loop.early_stop("validation_loss", mode: :min)
  |> Axon.Loop.run(train_pipeline, %{}, epochs: 100, compiler: EXLA)

cnn_model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(test_pipeline, cnn_trained_model_state, compiler: EXLA)

Chapter 8

defmodule CatsAndDogs do
  def pipeline(paths, batch_size, target_height, target_width) do
    paths
    |> Enum.shuffle()
    |> Task.async_stream(&amp;parse_image/1)
    |> Stream.filter(fn
      {:ok, {%StbImage{}, _}} -> true
      _ -> false
    end)
    |> Stream.map(&amp;to_tensors(&amp;1, target_height, target_width))
    |> Stream.chunk_every(batch_size, batch_size, :discard)
    |> Stream.map(fn chunks ->
      {img_chunk, label_chunk} = Enum.unzip(chunks)
      {Nx.stack(img_chunk), Nx.stack(label_chunk)}
    end)
  end
  
  def pipeline_with_augs(paths, batch_size, target_height, target_width) do
    paths
    |> Enum.shuffle()
    |> Task.async_stream(&amp;parse_image/1)
    |> Stream.filter(fn
      {:ok, {%StbImage{}, _}} -> true
      _ -> false
    end)
    |> Stream.map(&amp;to_tensors(&amp;1, target_height, target_width))
    |> Stream.map(&amp;random_flip(&amp;1, :height))
    |> Stream.map(&amp;random_flip(&amp;1, :width))
    |> Stream.chunk_every(batch_size, batch_size, :discard)
    |> Stream.map(fn chunks ->
      {img_chunk, label_chunk} = Enum.unzip(chunks)
      {Nx.stack(img_chunk), Nx.stack(label_chunk)}
    end)
  end

  defp random_flip({image, label}, axis) do
    if :rand.uniform() < 0.5 do
      {Nx.reverse(image, axes: [axis]), label}
    else
      {image, label}
    end
  end

  defp parse_image(path) do
    base = Path.basename(path)
    label = if String.contains?(base, "cat"), do: 0, else: 1

    case StbImage.read_file(path) do
      {:ok, img} -> {img, label}
      _error -> :error
    end
  end

  defp to_tensors({:ok, {img, label}}, target_height, target_width) do
    img_tensor =
      img
      |> StbImage.resize(target_height, target_width)
      |> StbImage.to_nx()
      |> Nx.divide(255)
      |> Nx.transpose(axes: [:channels, :height, :width])

    label_tensor = Nx.tensor([label])
    {img_tensor, label_tensor}
  end
end

batch_size = 32
target_height = 160
target_width = 160

{test_paths, train_paths} =
  Path.wildcard("train/cats_dogs/*.jpg")
  |> Enum.shuffle()
  |> Enum.split(1000)

{test_paths, val_paths} = test_paths |> Enum.split(750)

train_pipeline =
  CatsAndDogs.pipeline_with_augs(
    train_paths,
    batch_size,
    target_height,
    target_width
  )

test_pipeline =
  CatsAndDogs.pipeline(
    test_paths,
    batch_size,
    target_height,
    target_width
  )

val_pipeline =
  CatsAndDogs.pipeline(
    val_paths,
    batch_size,
    target_height,
    target_width
  )

Enum.take(train_pipeline, 1)

{cnn_base, cnn_base_params} = AxonOnnx.import(
  "train/mobilenetv2-7.onnx", batch_size: batch_size
)

input_template = Nx.template({1, 3, target_height, target_width}, :f32)
Axon.Display.as_graph(cnn_base, input_template)

### Extract the original classification head
{_popped, cnn_base} = cnn_base |> Axon.pop_node()
{_popped, cnn_base} = cnn_base |> Axon.pop_node()

Wrap convolutional base into its own namespace

cnn_base = cnn_base |> Axon.namespace("feature_extractor")

Freeze the convolutional base so that we don’t use it for training

cnn_base = cnn_base |> Axon.freeze()

Flatten the features or use a global pooling layer. Also add some regularization via dropout

model =
  cnn_base
  |> Axon.global_avg_pool(channels: :first)
  |> Axon.dropout(rate: 0.2)
  |> Axon.dense(1)

Create training loop

loss = &amp;Axon.Losses.binary_cross_entropy(&amp;1, &amp;2,
  reduction: :mean,
  from_logits: true
)
optimizer = Polaris.Optimizers.adam(learning_rate: 1.0e-4)

trained_model_state =
  model
  |> Axon.Loop.trainer(loss, optimizer)
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.validate(model, val_pipeline)
  |> Axon.Loop.early_stop("validation_loss", mode: :min, patience: 5)
  |> Axon.Loop.run(
    train_pipeline,
    %{"feature_extractor" => cnn_base_params},
    epochs: 1,
    compiler: EXLA
  )

eval_model = model |> Axon.sigmoid()

eval_model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(test_pipeline, trained_model_state, compiler: EXLA)

Fine-tuning

model = model |> Axon.unfreeze(up: 50)

loss = &amp;Axon.Losses.binary_cross_entropy(&amp;1, &amp;2,
  reduction: :mean,
  from_logits: true)

optimizer = Polaris.Optimizers.rmsprop(learning_rate: 1.0e-5)

trained_model_state =
  model
  |> Axon.Loop.trainer(loss, optimizer)
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.validate(model, val_pipeline)
  |> Axon.Loop.early_stop("validation_loss", mode: :min, patience: 5)
  |> Axon.Loop.run(
    train_pipeline,
    trained_model_state,
    epochs: 1,
    compiler: EXLA
  )

eval_model = model |> Axon.sigmoid()

eval_model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(test_pipeline, trained_model_state, compiler: EXLA)

Back