From 7b9dd686fb9a05d9fce9a2e0bddf21bfdb4abcd1 Mon Sep 17 00:00:00 2001 From: Shubham Gupta Date: Tue, 7 Dec 2021 11:44:10 +0800 Subject: [PATCH 1/3] Add example for Kuzushiji MNIST dataset --- examples/vision/kmnist.exs | 110 +++++++++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 examples/vision/kmnist.exs diff --git a/examples/vision/kmnist.exs b/examples/vision/kmnist.exs new file mode 100644 index 000000000..044d519fd --- /dev/null +++ b/examples/vision/kmnist.exs @@ -0,0 +1,110 @@ +# Title: Kuzushiji MNIST + +# ── Section ── + +Mix.install([ + {:axon, "~> 0.1.0-dev", github: "elixir-nx/axon"}, + {:exla, "~> 0.1.0-dev", github: "elixir-nx/nx", sparse: "exla"}, + {:nx, "~> 0.1.0-dev", github: "elixir-nx/nx", sparse: "nx", override: true}, + {:scidata, "~> 0.1.3"} +]) + +# Configure default platform with accelerator precedence as tpu > cuda > rocm > host +EXLA.set_preferred_defn_options([:tpu, :cuda, :rocm, :host]) + +defmodule KuzushijiMnist do + require Axon + + alias Axon.Loop.State + + defp transform_images({bin, type, shape}) do + bin + |> Nx.from_binary(type) + |> Nx.reshape({elem(shape, 0), 784}) + |> Nx.divide(255.0) + |> Nx.to_batched_list(32) + # Test split + |> Enum.split(1750) + end + + defp transform_labels({bin, type, _}) do + bin + |> Nx.from_binary(type) + |> Nx.new_axis(-1) + |> Nx.equal(Nx.tensor(Enum.to_list(0..9))) + |> Nx.to_batched_list(32) + # Test split + |> Enum.split(1750) + end + + defp build_model(input_shape) do + Axon.input(input_shape) + |> Axon.dense(128, activation: :relu) + |> Axon.dropout() + |> Axon.dense(10, activation: :softmax) + end + + defp log_metrics( + %State{epoch: epoch, iteration: iter, metrics: metrics, step_state: pstate} = state, + mode + ) do + loss = + case mode do + :train -> + %{loss: loss} = pstate + "Loss: #{:io_lib.format('~.5f', [Nx.to_scalar(loss)])}" + + :test -> + "" + end + + metrics = + metrics + |> Enum.map(fn {k, v} -> "#{k}: #{:io_lib.format('~.5f', [Nx.to_scalar(v)])}" end) + |> Enum.join(" ") + + IO.write("\rEpoch: #{Nx.to_scalar(epoch)}, Batch: #{Nx.to_scalar(iter)}, #{loss} #{metrics}") + + {:continue, state} + end + + defp train_model(model, train_images, train_labels, epochs) do + model + |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005)) + |> Axon.Loop.metric(:accuracy, "Accuracy") + |> Axon.Loop.handle(:iteration_completed, &log_metrics(&1, :train), every: 50) + |> Axon.Loop.run(Stream.zip(train_images, train_labels), epochs: epochs, compiler: EXLA) + end + + defp test_model(model, model_state, test_images, test_labels) do + model + |> Axon.Loop.evaluator(model_state) + |> Axon.Loop.metric(:accuracy, "Accuracy") + |> Axon.Loop.handle(:iteration_completed, &log_metrics(&1, :test), every: 50) + |> Axon.Loop.run(Stream.zip(test_images, test_labels), compiler: EXLA) + end + + def run do + {images, labels} = Scidata.KuzushijiMNIST.download() + + {train_images, test_images} = transform_images(images) + {train_labels, test_labels} = transform_labels(labels) + + model = build_model({nil, 784}) |> IO.inspect() + + IO.write("\n\n Training Model \n\n") + + model_state = + model + |> train_model(train_images, train_labels, 5) + + IO.write("\n\n Testing Model \n\n") + + model + |> test_model(model_state, test_images, test_labels) + + IO.write("\n\n") + end +end + +KuzushijiMnist.run() From e0321d73894c07dc1e634e644093d75ae52e992f Mon Sep 17 00:00:00 2001 From: Shubham Gupta Date: Tue, 7 Dec 2021 11:49:11 +0800 Subject: [PATCH 2/3] Remove livebook headers --- examples/vision/kmnist.exs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/examples/vision/kmnist.exs b/examples/vision/kmnist.exs index 044d519fd..ebcb787aa 100644 --- a/examples/vision/kmnist.exs +++ b/examples/vision/kmnist.exs @@ -1,7 +1,3 @@ -# Title: Kuzushiji MNIST - -# ── Section ── - Mix.install([ {:axon, "~> 0.1.0-dev", github: "elixir-nx/axon"}, {:exla, "~> 0.1.0-dev", github: "elixir-nx/nx", sparse: "exla"}, From 1dd0147c08cc719f8bbb50452988328fb5de8a67 Mon Sep 17 00:00:00 2001 From: Shubham Gupta Date: Thu, 16 Dec 2021 13:46:47 +0800 Subject: [PATCH 3/3] Add incomplete implementation for metric learning --- examples/vision/kmnist.exs | 106 --------------- examples/vision/kmnist_metric_learning.exs | 149 +++++++++++++++++++++ 2 files changed, 149 insertions(+), 106 deletions(-) delete mode 100644 examples/vision/kmnist.exs create mode 100644 examples/vision/kmnist_metric_learning.exs diff --git a/examples/vision/kmnist.exs b/examples/vision/kmnist.exs deleted file mode 100644 index ebcb787aa..000000000 --- a/examples/vision/kmnist.exs +++ /dev/null @@ -1,106 +0,0 @@ -Mix.install([ - {:axon, "~> 0.1.0-dev", github: "elixir-nx/axon"}, - {:exla, "~> 0.1.0-dev", github: "elixir-nx/nx", sparse: "exla"}, - {:nx, "~> 0.1.0-dev", github: "elixir-nx/nx", sparse: "nx", override: true}, - {:scidata, "~> 0.1.3"} -]) - -# Configure default platform with accelerator precedence as tpu > cuda > rocm > host -EXLA.set_preferred_defn_options([:tpu, :cuda, :rocm, :host]) - -defmodule KuzushijiMnist do - require Axon - - alias Axon.Loop.State - - defp transform_images({bin, type, shape}) do - bin - |> Nx.from_binary(type) - |> Nx.reshape({elem(shape, 0), 784}) - |> Nx.divide(255.0) - |> Nx.to_batched_list(32) - # Test split - |> Enum.split(1750) - end - - defp transform_labels({bin, type, _}) do - bin - |> Nx.from_binary(type) - |> Nx.new_axis(-1) - |> Nx.equal(Nx.tensor(Enum.to_list(0..9))) - |> Nx.to_batched_list(32) - # Test split - |> Enum.split(1750) - end - - defp build_model(input_shape) do - Axon.input(input_shape) - |> Axon.dense(128, activation: :relu) - |> Axon.dropout() - |> Axon.dense(10, activation: :softmax) - end - - defp log_metrics( - %State{epoch: epoch, iteration: iter, metrics: metrics, step_state: pstate} = state, - mode - ) do - loss = - case mode do - :train -> - %{loss: loss} = pstate - "Loss: #{:io_lib.format('~.5f', [Nx.to_scalar(loss)])}" - - :test -> - "" - end - - metrics = - metrics - |> Enum.map(fn {k, v} -> "#{k}: #{:io_lib.format('~.5f', [Nx.to_scalar(v)])}" end) - |> Enum.join(" ") - - IO.write("\rEpoch: #{Nx.to_scalar(epoch)}, Batch: #{Nx.to_scalar(iter)}, #{loss} #{metrics}") - - {:continue, state} - end - - defp train_model(model, train_images, train_labels, epochs) do - model - |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005)) - |> Axon.Loop.metric(:accuracy, "Accuracy") - |> Axon.Loop.handle(:iteration_completed, &log_metrics(&1, :train), every: 50) - |> Axon.Loop.run(Stream.zip(train_images, train_labels), epochs: epochs, compiler: EXLA) - end - - defp test_model(model, model_state, test_images, test_labels) do - model - |> Axon.Loop.evaluator(model_state) - |> Axon.Loop.metric(:accuracy, "Accuracy") - |> Axon.Loop.handle(:iteration_completed, &log_metrics(&1, :test), every: 50) - |> Axon.Loop.run(Stream.zip(test_images, test_labels), compiler: EXLA) - end - - def run do - {images, labels} = Scidata.KuzushijiMNIST.download() - - {train_images, test_images} = transform_images(images) - {train_labels, test_labels} = transform_labels(labels) - - model = build_model({nil, 784}) |> IO.inspect() - - IO.write("\n\n Training Model \n\n") - - model_state = - model - |> train_model(train_images, train_labels, 5) - - IO.write("\n\n Testing Model \n\n") - - model - |> test_model(model_state, test_images, test_labels) - - IO.write("\n\n") - end -end - -KuzushijiMnist.run() diff --git a/examples/vision/kmnist_metric_learning.exs b/examples/vision/kmnist_metric_learning.exs new file mode 100644 index 000000000..3d995e1ee --- /dev/null +++ b/examples/vision/kmnist_metric_learning.exs @@ -0,0 +1,149 @@ +# Title: KMNIST - Metric Learning + +# ── Section ── + +Mix.install([ + {:axon, github: "elixir-nx/axon"}, + {:exla, github: "elixir-nx/nx", sparse: "exla"}, + {:nx, github: "elixir-nx/nx", sparse: "nx", override: true}, + {:scidata, "~> 0.1.3"} +]) + +{train_images, train_labels} = Scidata.KuzushijiMNIST.download() +{test_images, test_labels} = Scidata.KuzushijiMNIST.download_test() + +{train_images_bin, train_images_type, train_images_shape} = train_images + +train_images_tensor = + train_images_bin + |> Nx.from_binary(train_images_type) + |> Nx.reshape(train_images_shape) + |> Nx.divide(255.0) + +{train_labels_bin, train_labels_type, train_labels_shape} = train_labels + +train_labels_tensor = + train_labels_bin + |> Nx.from_binary(train_labels_type) + |> Nx.reshape(train_labels_shape) + +train_images_tensor[0] |> Nx.to_heatmap() + +# Create a mapping between class labels and the images in each label. This will be used to sample and create the training dataset. + +{test_images_bin, test_images_type, test_images_shape} = test_images + +test_images_tensor = + test_images_bin + |> Nx.from_binary(test_images_type) + |> Nx.reshape(test_images_shape) + |> Nx.divide(255.0) + +{test_labels_bin, test_labels_type, test_labels_shape} = test_labels + +test_labels_tensor = + test_labels_bin + |> Nx.from_binary(test_labels_type) + |> Nx.reshape(test_labels_shape) + +class_idx_to_train_idxs = + train_labels_tensor + |> Nx.to_flat_list() + |> Enum.with_index() + |> Enum.reduce(%{}, fn {record, index}, acc -> + old_records = Map.get(acc, record, []) + Map.put(acc, record, [index | old_records]) + end) + +class_idx_to_test_idxs = + test_labels_tensor + |> Nx.to_flat_list() + |> Enum.with_index() + |> Enum.reduce(%{}, fn {record, index}, acc -> + old_records = Map.get(acc, record, []) + Map.put(acc, record, [index | old_records]) + end) + +defmodule KuzushijiMNISTMetricLearning do + require Axon + alias Axon.Loop.State + import Nx.Defn + + defp transform_images({bin, type, shape}) do + bin + |> Nx.from_binary(type) + |> Nx.reshape({elem(shape, 0), 1, 28, 28}) + |> Nx.divide(255.0) + |> Nx.to_batched_list(32) + end + + defp l2_normalize(input_tensor) do + den = + Nx.multiply(input_tensor, input_tensor) |> Nx.sum(axes: [-1], keep_axes: true) |> Nx.sqrt() + + Nx.divide(input_tensor, den) + end + + def build_model(input_shape) do + Axon.input({nil, 28, 28, 3}) + |> Axon.conv(32, kernel_size: {3, 3}, strides: 2, activation: :relu) + |> Axon.conv(64, kernel_size: {3, 3}, strides: 2, activation: :relu) + |> Axon.conv(128, kernel_size: {3, 3}, strides: 2, activation: :relu) + |> Axon.global_avg_pool() + |> Axon.dense(8) + + # |> l2_normalize() + end + + defnp running_average(avg, obs, i) do + avg + |> Nx.multiply(i) + |> Nx.add(obs) + |> Nx.divide(Nx.add(i, 1)) + end + + defn init(model, init_optim) do + params = Axon.init(model) + + %{ + iteration: Nx.tensor(0), + model_state: params, + optimizer_state: init_optim.(params), + loss: Nx.tensor(0.0) + } + end + + defn batch_step(model, optim, real_images, state) do + iter = state[:iteration] + params = state[:model_state] + IO.puts(iter) + # Add code to compute cosine similarity for metric learning + end + + defp train_loop(model) do + {init_optim, optim} = Axon.Optimizers.adam(2.0e-3, b1: 0.5) + + step = &batch_step(model, optim, &1, &2) + init = fn -> init(model, init_optim) end + + Axon.Loop.loop(step, init) + end + + defp log_iteration(state) do + %State{epoch: epoch, iteration: iter, step_state: pstate} = state + + loss = "Loss: #{:io_lib.format('~.5f', [Nx.to_scalar(pstate[:loss])])}" + + "\rEpoch: #{Nx.to_scalar(epoch)}, batch: #{Nx.to_scalar(iter)} #{loss}" + end + + def run(train_tensor, test_tensor) do + model = build_model({nil, 28, 28, 3}) + + train_loop(model) + |> Axon.Loop.log(:iteration_completed, &log_iteration/1, :stdio, every: 50) + |> Axon.Loop.run(train_tensor, epochs: 10, compiler: EXLA) + end +end + +model = KuzushijiMNISTMetricLearning.run(train_images_tensor, test_images_tensor)