refactor(examples): attempt to improve clarity of where some hard-coded numbers coming from or mean (#512)

Author: grzuy

examples/basics/multi_input_example.exs

@@ -7,6 +7,8 @@ Mix.install([
 defmodule XOR do
   require Axon
 
+  @batch_size 32
+
   defp build_model(input_shape1, input_shape2) do
     inp1 = Axon.input("x1", shape: input_shape1)
     inp2 = Axon.input("x2", shape: input_shape2)
@@ -18,8 +20,8 @@ defmodule XOR do
   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)])
+    x1 = Nx.tensor(for _ <- 1..@batch_size, do: [Enum.random(0..1)])
+    x2 = Nx.tensor(for _ <- 1..@batch_size, do: [Enum.random(0..1)])
     y = Nx.logical_xor(x1, x2)
     {%{"x1" => x1, "x2" => x2}, y}
   end

examples/basics/multi_output_example.exs

@@ -7,6 +7,8 @@ Mix.install([
 defmodule Power do
   require Axon
 
+  @batch_size 32
+
   defp build_model do
     fc =
       Axon.input("input", shape: {nil, 1})
@@ -34,8 +36,7 @@ defmodule Power do
       Stream.unfold(
         Nx.Random.key(:erlang.system_time()),
         fn key ->
-          # Batch size of 32
-          {x, next_key} = Nx.Random.uniform(key, -10, 10, shape: {32, 1}, type: {:f, 32})
+          {x, next_key} = Nx.Random.uniform(key, -10, 10, shape: {@batch_size, 1}, type: {:f, 32})
           {{x, {Nx.pow(x, 2), Nx.pow(x, 3)}}, next_key}
         end
       )

examples/generative/fashionmnist_autoencoder.exs

@@ -8,7 +8,15 @@ Mix.install([
 defmodule FashionMNIST do
   require Axon
 
+  @batch_size 32
+  @image_channels 1
+  @image_side_pixels 28
+  @channel_value_max 255
+
   defmodule Autoencoder do
+    @image_channels 1
+    @image_side_pixels 28
+
     defp encoder(x, latent_dim) do
       x
       |> Axon.flatten()
@@ -17,8 +25,8 @@ defmodule FashionMNIST do
 
     defp decoder(x) do
       x
-      |> Axon.dense(784, activation: :sigmoid)
-      |> Axon.reshape({:batch, 1, 28, 28})
+      |> Axon.dense(@image_side_pixels**2, activation: :sigmoid)
+      |> Axon.reshape({:batch, @image_channels, @image_side_pixels, @image_side_pixels})
     end
 
     def build_model(input_shape, latent_dim) do
@@ -31,9 +39,9 @@ defmodule FashionMNIST do
   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(32)
+    |> Nx.reshape({elem(shape, 0), @image_channels, @image_side_pixels, @image_side_pixels})
+    |> Nx.divide(@channel_value_max)
+    |> Nx.to_batched(@batch_size)
   end
 
   defp train_model(model, train_images, epochs) do
@@ -48,15 +56,15 @@ defmodule FashionMNIST do
 
     train_images = transform_images(images)
 
-    model = Autoencoder.build_model({nil, 1, 28, 28}, 64) |> IO.inspect()
+    model = Autoencoder.build_model({nil, @image_channels, @image_side_pixels, @image_side_pixels}, 64) |> IO.inspect()
 
     model_state = train_model(model, train_images, 5)
 
     sample_image =
       train_images
       |> Enum.fetch!(0)
       |> Nx.slice_along_axis(0, 1)
-      |> Nx.reshape({1, 1, 28, 28})
+      |> Nx.reshape({1, @image_channels, @image_side_pixels, @image_side_pixels})
 
     sample_image |> Nx.to_heatmap() |> IO.inspect()
 

examples/generative/mnist_gan.exs

@@ -13,12 +13,17 @@ defmodule MNISTGAN do
   alias Axon.Loop.State
   import Nx.Defn
 
+  @batch_size 32
+  @image_channels 1
+  @image_side_pixels 28
+  @channel_value_max 255
+
   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(32)
+    |> Nx.reshape({elem(shape, 0), @image_channels, @image_side_pixels, @image_side_pixels})
+    |> Nx.divide(@channel_value_max)
+    |> Nx.to_batched(@batch_size)
   end
 
   defp build_generator(z_dim) do
@@ -32,9 +37,9 @@ defmodule MNISTGAN do
     |> Axon.dense(1024)
     |> Axon.leaky_relu(alpha: 0.9)
     |> Axon.batch_norm()
-    |> Axon.dense(784)
+    |> Axon.dense(@image_side_pixels**2)
     |> Axon.tanh()
-    |> Axon.reshape({:batch, 28, 28, 1})
+    |> Axon.reshape({:batch, @image_side_pixels, @image_side_pixels, @image_channels})
   end
 
   defp build_discriminator(input_shape) do
@@ -80,9 +85,9 @@ defmodule MNISTGAN do
     g_params = state[:generator][:model_state]
 
     # Update D
-    fake_labels = Nx.iota({32, 2}, axis: 1)
+    fake_labels = Nx.iota({@batch_size, 2}, axis: 1)
     real_labels = Nx.reverse(fake_labels)
-    {noise, random_next_key} = Nx.Random.normal(state[:random_key], shape: {32, 100})
+    {noise, random_next_key} = Nx.Random.normal(state[:random_key], shape: {@batch_size, 100})
 
     {d_loss, d_grads} =
       value_and_grad(d_params, fn params ->
@@ -162,7 +167,7 @@ defmodule MNISTGAN do
     preds = Axon.predict(model, pstate[:generator][:model_state], noise)
 
     preds
-    |> Nx.reshape({batch_size, 28, 28})
+    |> Nx.reshape({batch_size, @image_side_pixels, @image_side_pixels})
     |> Nx.to_heatmap()
     |> IO.inspect()
 
@@ -174,7 +179,7 @@ defmodule MNISTGAN do
     train_images = transform_images(images)
 
     generator = build_generator(100)
-    discriminator = build_discriminator({nil, 28, 28, 1})
+    discriminator = build_discriminator({nil, @image_side_pixels, @image_side_pixels, @image_channels})
 
     discriminator
     |> train_loop(generator)

examples/vision/cifar10.exs

@@ -8,23 +8,27 @@ Mix.install([
 defmodule Cifar do
   require Axon
 
+  @batch_size 32
+  @channel_value_max 255
+  @label_values Enum.to_list(0..9)
+
   defp transform_images({bin, type, shape}) do
     bin
     |> Nx.from_binary(type)
     |> Nx.reshape(shape, names: [:count, :channels, :width, :height])
     # Move channels to last position to match what conv layer expects
     |> Nx.transpose(axes: [:count, :width, :height, :channels])
-    |> Nx.divide(255.0)
-    |> Nx.to_batched(32)
+    |> Nx.divide(@channel_value_max)
+    |> Nx.to_batched(@batch_size)
     |> Enum.split(1500)
   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(32)
+    |> Nx.equal(Nx.tensor(@label_values))
+    |> Nx.to_batched(@batch_size)
     |> Enum.split(1500)
   end
 
@@ -39,7 +43,7 @@ defmodule Cifar do
     |> Axon.flatten()
     |> Axon.dense(64, activation: :relu)
     |> Axon.dropout(rate: 0.5)
-    |> Axon.dense(10, activation: :softmax)
+    |> Axon.dense(length(@label_values), activation: :softmax)
   end
 
   defp train_model(model, train_images, train_labels, epochs) do

examples/vision/cnn_image_denoising.exs

@@ -12,6 +12,9 @@ defmodule MnistDenoising do
   @noise_factor 0.4
   @batch_size 32
   @epochs 25
+  @image_channels 1
+  @image_side_pixels 28
+  @channel_value_max 255
 
   def run do
     {images, _} = Scidata.MNIST.download()
@@ -26,7 +29,7 @@ defmodule MnistDenoising do
     noisy_train_images |> Enum.take(1) |> hd() |> display_image()
 
     # Train with noisy images as input and train images as targets
-    model = build_model({nil, 1, 28, 28})
+    model = build_model({nil, @image_channels, @image_side_pixels, @image_side_pixels})
 
     model_state =
       model
@@ -46,8 +49,8 @@ defmodule MnistDenoising do
   defp transform_images({bin, type, shape}) do
     bin
     |> Nx.from_binary(type)
-    |> Nx.reshape({elem(shape, 0), 28, 28, 1})
-    |> Nx.divide(255.0)
+    |> Nx.reshape({elem(shape, 0), @image_side_pixels, @image_side_pixels, @image_channels})
+    |> Nx.divide(@channel_value_max)
     |> Nx.to_batched_list(@batch_size)
     # Test split
     |> Enum.split(1750)
@@ -63,7 +66,7 @@ defmodule MnistDenoising do
   defp display_image(images) do
     images
     |> Nx.slice_along_axis(0, 1)
-    |> Nx.reshape({28, 28, 1})
+    |> Nx.reshape({@image_side_pixels, @image_side_pixels, @image_channels})
     |> Nx.to_heatmap()
     |> IO.inspect()
   end

examples/vision/horses_or_humans.exs

@@ -15,10 +15,14 @@ defmodule HorsesOrHumans do
   # or you can use Req to download and extract the zip file and iterate
   # over the resulting data
   @directories "examples/vision/{horses,humans}/*"
+  @batch_size 32
+  @image_channels 4
+  @image_side_pixels 300
+  @channel_value_max 255
 
   def data() do
     Path.wildcard(@directories)
-    |> Stream.chunk_every(32, 32, :discard)
+    |> Stream.chunk_every(@batch_size, @batch_size, :discard)
     |> Task.async_stream(fn batch ->
       {inp, labels} = batch |> Enum.map(&parse_png/1) |> Enum.unzip()
       {Nx.stack(inp), Nx.stack(labels)}
@@ -29,7 +33,7 @@ defmodule HorsesOrHumans do
 
   defnp augment(inp) do
     # Normalize
-    inp = inp / 255.0
+    inp = inp / @channel_value_max
 
     # For now just a random flip
     if Nx.random_uniform({}) > 0.5 do
@@ -78,7 +82,7 @@ defmodule HorsesOrHumans do
   end
 
   def run() do
-    model = build_model({nil, 300, 300, 4}) |> IO.inspect()
+    model = build_model({nil, @image_side_pixels, @image_side_pixels, @image_channels}) |> IO.inspect()
     optimizer = Polaris.Optimizers.adam(learning_rate: 1.0e-4)
     centralized_optimizer = Polaris.Updates.compose(Polaris.Updates.centralize(), optimizer)
 

examples/vision/mnist.exs

@@ -9,12 +9,17 @@ Mix.install([
 defmodule Mnist do
   require Axon
 
+  @batch_size 32
+  @image_side_pixels 28
+  @channel_value_max 255
+  @label_values Enum.to_list(0..9)
+
   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(32)
+    |> Nx.reshape({elem(shape, 0), @image_side_pixels**2})
+    |> Nx.divide(@channel_value_max)
+    |> Nx.to_batched(@batch_size)
     # Test split
     |> Enum.split(1750)
   end
@@ -23,8 +28,8 @@ defmodule Mnist do
     bin
     |> Nx.from_binary(type)
     |> Nx.new_axis(-1)
-    |> Nx.equal(Nx.tensor(Enum.to_list(0..9)))
-    |> Nx.to_batched(32)
+    |> Nx.equal(Nx.tensor(@label_values))
+    |> Nx.to_batched(@batch_size)
     # Test split
     |> Enum.split(1750)
   end
@@ -33,7 +38,7 @@ defmodule Mnist do
     Axon.input("input", shape: input_shape)
     |> Axon.dense(128, activation: :relu)
     |> Axon.dropout()
-    |> Axon.dense(10, activation: :softmax)
+    |> Axon.dense(length(@label_values), activation: :softmax)
   end
 
   defp train_model(model, train_images, train_labels, epochs) do
@@ -56,7 +61,7 @@ defmodule Mnist do
     {train_images, test_images} = transform_images(images)
     {train_labels, test_labels} = transform_labels(labels)
 
-    model = build_model({nil, 784}) |> IO.inspect()
+    model = build_model({nil, @image_side_pixels**2}) |> IO.inspect()
 
     IO.write("\n\n Training Model \n\n")