Mix.install([
{:nx, "~> 0.12.1"},
{:kino_zoetrope, "~> 0.25.0"},
{:kino_rewind, "~> 0.2.0"},
{:csv, "~> 3.2"},
{:pythonx, "~> 0.4.2"},
{:kino_pythonx, "~> 0.1.0"},
{:nx_signal, "~> 0.3.0"},
{:exla, "~> 0.12.0"},
])Configure Nx to calculate on the GPU (if available) or hardware accelerated on the CPU.
Nx.global_default_backend(EXLA.Backend)The MNIST dataset contains a collection of handwritten digits (0-9) stored as 28x28 grayscale images with 255 values (8bit) per pixel and an additional label (0-9)
The data set is stored in a csv file with 785 (28*28+1) columns. We the the first 100 rows and plot them as 28x28 bitmap and 1x1 bitmap.
You can download it here: https://www.kaggle.com/api/v1/datasets/download/oddrationale/mnist-in-csv
csv_stream =
Kino.FS.file_path("mnist_train.csv")
|> File.stream!()
|> Stream.drop(1)
|> CSV.decode!()
# per performance reasons we process only 30 items of the data set
batch_size = 30
data =
csv_stream
# read only the bitmap, drop the first columnn (label)
|> Stream.map(&Enum.drop(Enum.map(&1, fn x -> String.to_integer(x) end), 1))
|> Enum.take(batch_size)
|> Enum.to_list()
|> Nx.tensor()
|> Nx.reshape({:auto, 1, 28, 28})
# convert 0..255 to 0..1.0
|> Nx.divide(255)
labels =
csv_stream
# read only the labels, keep only the first column
|> Stream.map(&Enum.take(Enum.map(&1, fn x -> String.to_integer(x) end), 1))
|> Enum.take(batch_size)
|> Enum.to_list()
|> Nx.tensor()
|> Nx.reshape({:auto, 1})
# One-Hot-Encoding is a way to encoded an intenger n as
# as a 1D vector of zeros with only the nth field set to 1
# it can be thought of as a discrete probability distribution (or histogram)
# saying 100% chance of the the number being equal to n
labels_one_hot =
Nx.iota({10})
|> Nx.equal(labels)
[
data |> Nx.reshape({:auto, 28, 28}),
labels |> Nx.reshape({:auto, 1, 1, 1}),
labels_one_hot |> Nx.reshape({:auto, 1, 10, 1})
]
|> KinoZoetrope.TensorStack.new(
titel: "Handwritten digits",
labels: ["Bitmap", "Label", "Labels (One Hot)"],
frame_label: "Input Image"
)kernelA =
Nx.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
|> Nx.new_axis(2)
|> Nx.new_axis(0)
kernelB =
Nx.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
|> Nx.transpose(axes: [1, 0])
|> Nx.new_axis(2)
|> Nx.new_axis(0)
kernelC =
Nx.tensor([[0.5, 1, 0.5], [1, 2, 1], [0.5, 1, 0.5]])
|> Nx.transpose(axes: [1, 0])
|> Nx.new_axis(2)
|> Nx.new_axis(0)
allKernels =
Nx.concatenate([
kernelA,
kernelB,
kernelC
])
[kernelA, kernelB, kernelC, allKernels]
|> KinoZoetrope.TensorStack.new(
titel: "4 Example Kernels",
size: 100,
labels: ["Horizontal", "Vertical", "blur", "All three kernels stacked"]
)kernelA = Nx.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]]) |> Nx.new_axis(0)
kernelB =
Nx.tensor([[1, 0, 1], [0, 0, 0], [-1, -2, -1]])
|> Nx.transpose(axes: [1, 0])
|> Nx.new_axis(0)
kernelC =
Nx.tensor([[0.5, 1, 0.5], [1, 2, 1], [0.5, 1, 0.5]])
|> Nx.transpose(axes: [1, 0])
|> Nx.new_axis(0)
bothKernels =
Nx.concatenate([
kernelA,
kernelB,
kernelC
])
Nx.shape(data)kernelA =
Nx.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
|> Nx.new_axis(0)
|> Nx.new_axis(0)
kernelB =
Nx.tensor([[1, 0, 1], [0, 0, 0], [-1, -2, -1]])
|> Nx.transpose(axes: [1, 0])
|> Nx.new_axis(0)
|> Nx.new_axis(0)
kernelC =
Nx.tensor([[0.5, 1, 0.5], [1, 2, 1], [0.5, 1, 0.5]])
|> Nx.transpose(axes: [1, 0])
|> Nx.new_axis(0)
|> Nx.new_axis(0)
bothKernels =
Nx.concatenate([
kernelA,
kernelB,
kernelC
])
# result1 = Nx.conv(data, bothKernels, padding: :same)
result1 =
data
|> Nx.vectorize(:batch)
|> Nx.pad(0, [{0, 0, 0}, {1, 1, 0}, {1, 1, 0}])
|> NxSignal.Convolution.convolve(
bothKernels
|> Nx.vectorize(:channel),
mode: :valid
)
|> Nx.devectorize()
|> Nx.squeeze(axes: [2])
[
data |> Nx.transpose(axes: [0, 2, 3, 1])
| for c <- 2..0//-1,
reduce: [
result1
|> Nx.reshape({:auto, 3, 28, 28})
|> Nx.transpose(axes: [0, 2, 3, 1])
|> Nx.multiply(255 / 4.0)
] do
l ->
[Nx.slice_along_axis(result1, c, 1, axis: 1) |> Nx.transpose(axes: [0, 2, 3, 1]) | l]
end
]
|> KinoZoetrope.TensorStack.new(
titel: "Result of Convolution with 4 Kernels",
cmap: :viridis,
labels: [
"Original Input",
"Horizontal Edges",
"Vertical Edges",
"Blurred",
"Hor/Ver/Blur as RGB"
],
frame_label: "Input Image"
)key = Nx.Random.key(0)
{sixteenKernels, new_key} =
Nx.Random.uniform(key, shape: {16, 1, 3, 3})
|> then(fn {t, k} -> {Nx.subtract(t, 0.5), k} end)
{sixteenBiases, new_key} =
Nx.Random.uniform(new_key, shape: {1, 16, 1, 1})
|> then(fn {t, k} -> {Nx.subtract(t, 0.5), k} end)
[
sixteenBiases |> Nx.transpose(axes: [1, 0, 2, 3]),
sixteenKernels |> Nx.transpose(axes: [0, 2, 3, 1])
]
|> KinoZoetrope.TensorStack.new(
titel: "16 Random kernels",
labels: ["Bias", "Weights"],
frame_label: "Kernel"
)# result_layer_1 = Nx.conv(data, sixteenKernels, padding: :same) |> Nx.add(sixteenBiases)
result_layer_1 =
data
|> Nx.pad(0, [{0, 0, 0}, {0, 0, 0}, {1, 1, 0}, {1, 1, 0}])
|> Nx.vectorize(:batch)
|> NxSignal.Convolution.convolve(
sixteenKernels
|> Nx.vectorize(:kernel),
mode: :valid
)
|> Nx.devectorize()
|> Nx.sum(axes: [2])
[
data |> Nx.transpose(axes: [0, 2, 3, 1])
| for c <- 0..15 do
Nx.slice_along_axis(result_layer_1, c, 1, axis: 1) |> Nx.transpose(axes: [0, 2, 3, 1])
end
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 1 Convolution results",
cmap: :viridis,
show_meta: false,
labels: ["Original Input" | for(c <- 1..16, do: "Result for kernel #{c}")],
frame_label: "Input Image"
)Next we discard all negative values by setting them to 0. This is done to do a non-linear operation (ie something other than a matrix multiplication, ie. something other than a weighted sum). This adds some spice. Chaining only linear operations together yield the same a result as a single linear operation.
result_layer_1_relu = Nx.max(0, result_layer_1)
[
data |> Nx.transpose(axes: [0, 2, 3, 1])
| for c <- 0..15 do
Nx.slice_along_axis(result_layer_1_relu, c, 1, axis: 1) |> Nx.transpose(axes: [0, 2, 3, 1])
end
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 1 Relu results (max(0, x))",
cmap: :viridis,
show_meta: false,
labels: ["Original Input" | for(c <- 1..16, do: "Result for kernel #{c} Relu")],
frame_label: "Input Image"
)Next we downsample the 28x28 images to 14x14 via maxpooling. Maxpooling chunks the images into 2x2 patches and keeps only the largest of its values.
result_layer_1_relu_max_pooled =
result_layer_1_relu |> Nx.window_max({1, 1, 2, 2}, strides: [1, 1, 2, 2])
[
data |> Nx.transpose(axes: [0, 2, 3, 1])
| for c <- 0..15 do
Nx.slice_along_axis(result_layer_1_relu_max_pooled, c, 1, axis: 1)
|> Nx.transpose(axes: [0, 2, 3, 1])
end
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 1 Pooling Results",
cmap: :viridis,
show_meta: false,
labels: ["Original Input" | for(c <- 1..16, do: "K#{c} result max-pooled")],
frame_label: "Input Image"
)Before we start with the next layer of calculations, lets define a helper function for cleaner plotting of our vectors.
The make_grid function will take a tensor of shape {b, c, w, h} and turn in into a vector of shape {b, 1, w*c/4, h*4}. It moves the values from the different channels into the spatial (width, height) dimensions for us to look at multiple channels at once.
defmodule Grid do
import Nx.Defn
# Help to reshape channels into a grid
defn make_grid(tensor, opts \\ []) do
opts = keyword!(opts, grid_size: 4)
grid_size = opts[:grid_size]
{batch, _channels, w, h} = Nx.shape(tensor)
tensor
|> Nx.reshape({batch, grid_size, :auto, w, h})
|> Nx.transpose(axes: [0, 1, 3, 2, 4])
|> Nx.reshape({batch, 1, grid_size * w, :auto})
end
end
Kino.nothing()Now by using our helper function we can look at the full output (14 * 14 * 16 = 56x56 values) of layer 1 for each input:
result_layer_1_relu_max_pooled_flatten =
result_layer_1_relu_max_pooled |> Nx.reshape({:auto})
[
data |> Nx.transpose(axes: [0, 2, 3, 1]),
result_layer_1_relu_max_pooled
|> Grid.make_grid(grid_size: 4)
|> Nx.transpose(axes: [0, 2, 3, 1])
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 1 Pooling flattend into 2d array",
labels: ["Original Input", "Result Layer 1 2D-Flattened"],
size: [200, 400],
cmap: :viridis,
frame_label: "Input Image"
)All the (14 * 14 * 16 = 56 * 56) values calculated by Layer 1 are passed to Layer 2. Layer 2 will use 32 more kernels to convolve the 16 channels produced by Layer 1.
# Generate 32 kernels of size 16x3x3 (3x3 across width*height and across all 16 channels)
key = Nx.Random.key(1)
{thirtyTwoKernels, new_key} =
Nx.Random.uniform(key, shape: {32, 16, 3, 3})
|> then(fn {t, k} -> {Nx.subtract(t, 0.5), k} end)
{thirtyTwoBiases, new_key} =
Nx.Random.uniform(new_key, shape: {1, 32, 1, 1})
|> then(fn {t, k} -> {Nx.subtract(t, 0.5), k} end)
Nx.shape(thirtyTwoKernels)Layer 2 only two times as many kernels as in Layer 2 but each kernel has 16 times as large. Lets plot all the kernels channels values:
[
thirtyTwoBiases |> Nx.transpose(axes: [1, 0, 2, 3])
| for c <- 0..15 do
Nx.slice_along_axis(thirtyTwoKernels, c, 1, axis: 1) |> Nx.transpose(axes: [0, 2, 3, 1])
end
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 2 uses 32 Random kernels for 16 channels (32*16*3*3 in total)",
show_meta: false,
size: 80,
legend: false,
labels: [
"Bias"
| for(
c <- 1..16,
do: "Kernel channel #{c}"
)
],
frame_label: "Kernel"
)Now lets apply the kernels to the output of Layer 1. The convolution produces tensors / images that are 14*14 pixels wide and high and have 32 channels per channel. (32 channels because we had 32 kernels, each producing one value)
Nx.shape(result_layer_1_relu_max_pooled)Nx.names(result_layer_1_relu_max_pooled)# result_layer_2 = Nx.conv(result_layer_1_relu_max_pooled, thirtyTwoKernels, padding: :same)
result_layer_2 =
result_layer_1_relu_max_pooled
|> Nx.pad(0, [{0, 0, 0}, {0, 0, 0}, {1, 1, 0}, {1, 1, 0}])
|> Nx.vectorize(:batch)
|> NxSignal.Convolution.convolve(
thirtyTwoKernels
|> Nx.vectorize(:channel),
mode: :valid
)
|> Nx.devectorize()
|> Nx.squeeze(axes: [2])
Nx.shape(result_layer_2)Lets inspect the result visually. At this point, for each initial input we have 32 intermediate resulting images to look at.
[
data |> Nx.transpose(axes: [0, 2, 3, 1])
| for c <- 0..31 do
Nx.slice_along_axis(result_layer_2, c, 1, axis: 1) |> Nx.transpose(axes: [0, 2, 3, 1])
end
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 2 Convolution results",
cmap: :viridis,
show_meta: false,
legend: false,
labels: ["Original Input" | for(c <- 1..32, do: "Result for kernel #{c}")],
frame_label: "Input Image"
)Rinse and repeat: discard negative values by setting them to 0. This is a per-pixel operation not chaning the number of pixels. And looking at the results again:
result_layer_2_relu = Nx.max(0, result_layer_2)
[
data |> Nx.transpose(axes: [0, 2, 3, 1])
| for c <- 0..31 do
Nx.slice_along_axis(result_layer_2_relu, c, 1, axis: 1) |> Nx.transpose(axes: [0, 2, 3, 1])
end
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 2 Relu Results (max(0, 1))",
cmap: :viridis,
show_meta: false,
legend: false,
labels: [
"Original Input"
| for(c <- 1..32, do: "Result K#{c} Relu")
],
frame_label: "Input Image"
)As in Layer 1 we apply some max-pooling again to reduce the number of pixels. From each 2x2 chunk of neighbors we only keep the largest value again.
Be plot the results again and can see that each channel has only 7x7 pixels.
result_layer_2_relu_max_pooled =
result_layer_2_relu |> Nx.window_max({1, 1, 2, 2}, strides: [1, 1, 2, 2])
[
data |> Nx.transpose(axes: [0, 2, 3, 1])
| for c <- 0..31 do
Nx.slice_along_axis(result_layer_2_relu_max_pooled, c, 1, axis: 1)
|> Nx.transpose(axes: [0, 2, 3, 1])
end
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 2 Pooling result",
cmap: :viridis,
show_meta: false,
legend: false,
labels: [
"Original Input"
| for(c <- 1..32, do: "Result K_#{c} Relu MaxPooled")
],
frame_label: "Input Image"
)The calculations of Layer 2 are done. Be can reshape all the resulting values into 2D grid to look at all of them at once:
[
data |> Nx.transpose(axes: [0, 2, 3, 1]),
result_layer_2_relu_max_pooled |> Grid.make_grid() |> Nx.transpose(axes: [0, 2, 3, 1])
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 2 Pooling flattend into 2D array",
labels: ["Original Input", "Result Layer 2 Flattened"],
cmap: :viridis,
size: [100, 400],
frame_label: "Input Image"
)The next layer (Layer 3) is a different kind of layer. It does not convolve the tensors with a kernel but just calculates a weighted some across all pixels (a vector dot product). It does so with 120 different list of weights. Each list has as many weights as the total number of input values.
From layer 2 we expect to get (7 * 7 * 32 = 1568) values. Layer 3 shall product 120 resulting values by weighting and summing the 1568 values in 100 different ways.
So Layer 3 needs (120 * 1568 = 188160) different coefficients to multiply with.
We first need to create tensor to contain these numbers. We can also plot this 2D matrix:
key = Nx.Random.key(2)
{randomMatrix_layer3, new_key} =
Nx.Random.uniform(key, shape: {32 * 7 * 7, 120})
|> then(fn {t, k} -> {Nx.subtract(t, 0.5), k} end)
{bias_layer3, new_key} =
Nx.Random.uniform(new_key, shape: {1, 120})
|> then(fn {t, k} -> {Nx.subtract(t, 0.5), k} end)
[
bias_layer3
|> Nx.reshape({1, 10, 12, 1}),
randomMatrix_layer3
|> Nx.reshape({1, 120, :auto, 1})
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 3 Weight Matrix",
size: [100, 800],
show_meta: true,
legend: false,
cmap: :viridis,
labels: ["Bias", "Matrix Values"]
)Next we need to multiply this matrix via matrix multiplication with the values we got from Layer 2. From the perspective of Layer 3 all the pixels That Layer 2 produces are just a very long 1D-vector of numbers (7 * 7 * 32).
The result of Layer 2 is then just a list of 100 numbers per initial input.
result_layer_2_relu_max_pooled_flatten =
result_layer_2_relu_max_pooled |> Nx.reshape({:auto, 7 * 7 * 32})
result_layer_3 =
result_layer_2_relu_max_pooled_flatten
|> Nx.dot(randomMatrix_layer3)
|> Nx.add(bias_layer3)
Nx.shape(result_layer_3)Lets plot the result so far together with the initial input and the values we got from Layer 2.
Layer 3 produces just a flat list of 100 numbers but when plotting them its nice to look at if we reshape them into a 10x10 grid.
[
data |> Nx.transpose(axes: [0, 2, 3, 1]),
result_layer_2_relu_max_pooled
|> Grid.make_grid()
|> Nx.transpose(axes: [0, 2, 3, 1]),
result_layer_3
|> Nx.reshape({:auto, 12, 10})
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 3 Matrix Multiplication",
show_meta: false,
labels: ["Original Input", "Input from Layer 2", "Result Layer 3 Matrix Multiplication"],
cmap: :viridis,
legend: false,
size: [150, 300, 250],
frame_label: "Input Image"
)Next we again discard all negative values by setting them to 0. And plot the result again, just for fun
result_layer_3_relu = Nx.max(0, result_layer_3)
[
data |> Nx.transpose(axes: [0, 2, 3, 1]),
result_layer_2_relu_max_pooled
|> Grid.make_grid()
|> Nx.transpose(axes: [0, 2, 3, 1]),
result_layer_3
|> Nx.reshape({:auto, 12, 10}),
result_layer_3_relu
|> Nx.reshape({:auto, 12, 10})
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 3 Relu Activation",
show_meta: false,
cmap: :viridis,
legend: [false, false, true, true],
size: [100, 400, 200, 200],
labels: ["Original Input", "Input from Layer 2", "Before Relu", "Result Layer 3 Relu"],
frame_label: "Input Image"
)Layer 4 does the same as Layer 3: Take all values as as flat list and multiply them by a bunch of weights to produce a new list of values.
We get 120 values from Layer 3 and want to produce only 10 different output values. This means we need 120 * 10 = 1200 weight-values, or a 120 * 10 matrix:
key = Nx.Random.key(2)
{randomMatrix_layer4, new_key} =
Nx.Random.uniform(key, shape: {120, 10})
|> then(fn {t, k} -> {Nx.subtract(t, 0.5), k} end)
{bias_layer4, new_key} =
Nx.Random.uniform(new_key, shape: {10})
|> then(fn {t, k} -> {Nx.subtract(t, 0.5), k} end)
[
bias_layer4
|> Nx.reshape({1, 10, 1, 1}),
randomMatrix_layer4
|> Nx.reshape({1, :auto, 120, 1})
]
|> KinoZoetrope.TensorStack.new(
titel: "Layer 4 Weight Matrix",
size: [10, 600],
show_meta: true,
cmap: :viridis,
legend: false,
labels: ["Bias", "Weights"]
)The matrix multiplication done by Layer 4 produces only 10 values for each initial input image.
result_layer_4 =
result_layer_3_relu
|> Nx.dot(randomMatrix_layer4)
|> Nx.add(bias_layer4)
Nx.shape(result_layer_4)These 10 values shall represent a probability distribution describing the probability of the input image containing the respective value between 0 and 9.
Since all the convolution kernels and weight matrix values were randomly generated we can expect the resulting 10-element-vector to also just containing 10 pretty random values.
[
data |> Nx.transpose(axes: [0, 2, 3, 1]),
result_layer_3_relu
|> Nx.reshape({:auto, 12, 10}),
result_layer_4
|> Nx.reshape({:auto, 10, 1, 1})
]
|> KinoZoetrope.TensorStack.new(
titel: "Output Layer",
show_meta: false,
cmap: :viridis,
size: [150, 150, 20],
labels: [
"Input",
"Input from Layer 3",
"Result Layer 4"
],
frame_label: "Input Image"
)Even if the output vector does not contain useful results yet, we still want to make sure that at least the vector formally represents a probability distribution summing to 1. Therefor we normalie the vector via softmax.
We also can use Nx.argmax to pick the index with the largest value from the vector.
softmax =
result_layer_4
|> Nx.exp()
|> then(fn t ->
t |> Nx.divide(Nx.sum(t, axes: [1], keep_axes: true))
end)
argmax =
result_layer_4
|> Nx.argmax(axis: 1)
softmax_log = softmax |> Nx.add(Nx.Constants.epsilon({:f, 32})) |> Nx.log()
[
data |> Nx.transpose(axes: [0, 2, 3, 1]),
result_layer_4
|> Nx.reshape({:auto, 10, 1, 1}),
argmax |> Nx.reshape({:auto, 1, 1, 1}),
softmax |> Nx.reshape({:auto, 10, 1, 1}),
softmax_log |> Nx.reshape({:auto, 10, 1, 1})
]
|> KinoZoetrope.TensorStack.new(
titel: "Output Layer",
size: [210, 20, 210, 20, 20],
show_meta: false,
cmap: :viridis,
labels: [
"Input",
"Result Layer 4",
"Result Layer 4 argmax",
"Result layer 4 softmax",
"Log(Softmax)"
],
frame_label: "Input Image"
)To actuall train our network we need a way to compare the output agains the expected value to measure the output quality. Is works kind of like test-driven-development. We need to write a unit-test, but one that can not only fail or pass, but tell how strongly it failed.
This can then be used to guide the computer to a better solution (ie. better values for our kernels and weight matrices)
This test-case is called lost function and this the most important part of the whole process. A good loss function will make the network learn. A bad loss function will lead to bad results that do not improve.
In our case the loss-function is a binary function that takes the expected label and the output of the network and produces a number telling the difference (0.0 means the output matches the expected label value, and large loss values mean that the output differs a lot from the actual label).
In our case the label (one-hot encoded 10-element-vector) and the output (also 10-element vector) can both be seen as probability distribution.
There are many ways to measure the difference between two distributions. A simple one is called cross-entropy. It can be written as -dot(log(a), log(b)), ie we negative sum of the pairwise product of elementwise logarithm. It produces a single non-negative value that is 0 if the vectors match and greater otherwise.
Lets to the calculation and plot everything together.
all_parameters =
[
{"Layer 1 weights", {16, 9}, sixteenKernels},
{"Layer 1 bias", {16, 1}, sixteenBiases},
{"Layer 2 weights", {32 * 2, 3 * 3 * 4 * 2}, thirtyTwoKernels},
{"Layer 2 bias", {8, 4}, thirtyTwoBiases},
{"Layer 3 weights", {120, :auto}, randomMatrix_layer3},
{"Layer 3 bias", {10, 12}, bias_layer3},
{"Layer 4 weights", {20, :auto}, randomMatrix_layer4},
{"Layer 4 bias", {10, 1}, bias_layer4}
]
cross_entropy =
Nx.multiply(
labels_one_hot
|> Nx.reshape({:auto, 1, 1, 1, 10}),
softmax_log
|> Nx.reshape({:auto, 1, 1, 1, 10})
)
|> Nx.sum(axes: [-1])
|> Nx.negate()
[
data |> Nx.transpose(axes: [0, 2, 3, 1]),
labels_one_hot |> Nx.reshape({:auto, 10, 1, 1}),
softmax |> Nx.reshape({:auto, 10, 1, 1}),
softmax_log |> Nx.reshape({:auto, 10, 1, 1}),
cross_entropy
|> Nx.reshape({1, :auto, 10, 1})
|> Nx.add(Nx.Constants.epsilon({:f, 32}))
|> Nx.log()
# all_parameters
# |> Enum.map(&Nx.flatten(elem(&1, 2)))
# |> Nx.concatenate()
# |> Nx.reshape({1, 674, :auto, 1})
]
|> Enum.concat(all_parameters |> Enum.map(fn {_, {a, b}, p} -> Nx.reshape(p, {1, a, b, 1}) end))
|> KinoZoetrope.TensorStack.new(
titel: "Random Matrix",
show_meta: false,
labels:
[
"Input",
"Expected as One Hot",
"Softmax(Output)",
"Log(Softmax(Output))",
"Cross Entropy (all Inputs)"
# "all parameters"
]
|> Enum.concat(
all_parameters
|> Enum.map(fn {l, _, p} -> "#{l} (total #{Nx.size(p)})" end)
),
legend: false,
cmap: :viridis,
size: [100, 20, 20, 20, 200, 80, 10, 200, 50, 500, 50, 100, 10],
markers: [
%{
attrs: %{
# SVG attributes
"stroke" => "magenta",
"stroke-width" => 2,
"fill" => "none",
"rx" => 10,
"ry" => 10
},
# indices of the images to add the markers to, here the first (0) and second (1) image
for: [4],
# list of {x,y} coordinates. The length if the list
# should match the number of frames of the image
points: 0..(Nx.size(labels) - 1) |> Enum.map(&{rem(&1, 10), div(&1, 10)})
}
],
frame_label: "Input Image"
)labels_one_hot |> Nx.shape()Now that we have unsterstood each step of the calculation we can but everything togher into a single module and see that the whole network can be expressed in merely 25 lines of code.
defmodule Handwriting do
import Nx.Defn
defn loss(
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
matrixA,
matABias,
matrixB,
matBBias,
data,
labels_one_hot
) do
output =
data
|> Nx.rename([:batch, :channel, :x, :y])
|> Nx.conv(sixteenKernels, padding: [{1, 1}, {1, 1}])
# |> Nx.vectorize(:batch)
# |> NxSignal.Convolution.convolve(
# sixteenKernels
# |> Nx.vectorize(:channel),
# mode: :valid
# )
# |> Nx.devectorize()
# |> Nx.squeeze(axes: [1])
|> Nx.rename([:batch, :channel, :x, :y])
|> Nx.add(sixteenBiases)
|> Nx.max(Nx.tensor(0))
|> Nx.window_max({1, 1, 2, 2}, strides: [1, 1, 2, 2])
|> Nx.conv(thirtyTwoKernels, padding: [{1, 1}, {1, 1}])
# |> Nx.pad(0, [{0, 0, 0}, {0, 0, 0}, {1, 1, 0}, {1, 1, 0}])
# |> Nx.vectorize(:batch)
# |> NxSignal.Convolution.convolve(
# thirtyTwoKernels
# |> Nx.vectorize(:channel),
# mode: :valid
# )
# |> Nx.devectorize()
# |> Nx.squeeze(axes: [1])
|> Nx.rename([:batch, :channel, :x, :y])
|> Nx.add(thirtyTwoBiases)
|> Nx.max(Nx.tensor(0))
|> Nx.window_max({1, 1, 2, 2}, strides: [1, 1, 2, 2])
|> Nx.reshape({:auto, 7 * 7 * 32})
|> Nx.dot(matrixA)
|> Nx.add(matABias)
|> Nx.max(Nx.tensor(0))
|> Nx.dot(matrixB)
|> Nx.add(matBBias)
output_exp = Nx.exp(output)
output_exp
|> Nx.divide(Nx.sum(output_exp, axes: [1], keep_axes: true))
|> Nx.add(Nx.Constants.epsilon({:f, 32}))
|> Nx.log()
|> Nx.multiply(labels_one_hot)
|> Nx.sum(axes: [-1])
|> Nx.negate()
end
defn mean_loss(
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
matrixA,
matABias,
matrixB,
matBBias,
data,
labels_one_hot
) do
loss(
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
matrixA,
matABias,
matrixB,
matBBias,
data,
labels_one_hot
)
|> Nx.mean()
end
end
Kino.nothing()We can use Nx.Defn to JIT compile the calculation:
params =
for v <- [
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
randomMatrix_layer3,
bias_layer3,
randomMatrix_layer4,
bias_layer4,
data,
labels_one_hot
],
reduce: [] do
tail ->
[Nx.template(Nx.shape(v), Nx.type(v)) | tail]
end
|> Enum.reverse()
compiled =
Nx.Defn.compile(
&Handwriting.loss/10,
params
)
compiled_mean =
Nx.Defn.compile(
&Handwriting.mean_loss/10,
params
)And then run it again on our dataset:
cross_entropy_all_at_once =
compiled.(
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
randomMatrix_layer3,
bias_layer3,
randomMatrix_layer4,
bias_layer4,
data,
labels_one_hot
)
Kino.nothing()And just for fun print the the results again, together with the input images and parameter values from all layers:
all_params_combined =
all_parameters
|> Enum.map(&Nx.flatten(elem(&1, 2)))
|> Nx.concatenate()
[
data |> Nx.transpose(axes: [0, 2, 3, 1]),
cross_entropy_all_at_once |> Nx.reshape({1, :auto, 10, 1}),
all_params_combined |> Nx.reshape({1, 10, :auto, 1})
]
|> KinoZoetrope.TensorStack.new(
titel: "Random Matrix",
show_meta: false,
size: [100, 100, 400],
labels: [
"Original Input",
"Cross Entropy (all Inputs)",
"all parameters combined (total #{Nx.size(all_params_combined)})"
],
legend: false,
markers: [
%{
attrs: %{
# SVG attributes
"stroke" => "magenta",
"stroke-width" => 2,
"fill" => "none",
"rx" => 10,
"ry" => 10
},
# indices of the images to add the markers to, here the first (0) and second (1) image
for: [1],
# list of {x,y} coordinates. The length if the list
# should match the number of frames of the image
points: 0..(Nx.size(labels) - 1) |> Enum.map(&{rem(&1, 10), div(&1, 10)})
}
],
frame_label: "Input image"
)For training the network to actually produce the expected results we we just need to take the partial derivative of the loss value, regarding all the parameter values. and then nudge its parameter a little bit down the the slope (ie. into the direction in which the loss is reduced most).
Nx.Defn.grad allows to to caculate the loss value again, but instead of actually giving us the resulting loss value it gives us all the slope value for all the parameters.
This multi dimensional derivative/slope is called gradient.
Notice that Nx.Defn.grad returns exactly as many results as it is given arguments.
{k16_grad, b16_grad, k32_grad, b32_grad, mat1_grad, mat1_bias_grad, mat2_grad, mat2_bias_grad} =
Nx.Defn.jit_apply(
fn d,
l,
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
randomMatrix_layer3,
bias_layer3,
randomMatrix_layer4,
bias_layer4 ->
Nx.Defn.grad(
{
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
randomMatrix_layer3,
bias_layer3,
randomMatrix_layer4,
bias_layer4
},
fn {kernel16, b16, kernel32, b32, matrixA, m1b, matrixB, m2b} ->
Handwriting.mean_loss(
kernel16,
b16,
kernel32,
b32,
matrixA,
m1b,
matrixB,
m2b,
d,
l
)
end
)
end,
[
data,
labels_one_hot,
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
randomMatrix_layer3,
bias_layer3,
randomMatrix_layer4,
bias_layer4
]
)
Kino.nothing()Next we can plot all the parameter valus again, together with their gradients (slope) values. Notice that there a as many gradient values as parameter values.
[
sixteenKernels
|> Nx.reshape({1, :auto, 3, 3})
|> Grid.make_grid(grid_size: 4)
|> Nx.transpose(axes: [0, 2, 3, 1]),
k16_grad
|> Nx.reshape({1, :auto, 3, 3})
|> Grid.make_grid(grid_size: 4)
|> Nx.transpose(axes: [0, 2, 3, 1]),
thirtyTwoKernels
|> Nx.reshape({1, :auto, 3, 3})
|> Grid.make_grid(grid_size: 16)
|> Nx.transpose(axes: [0, 2, 3, 1]),
k32_grad
|> Nx.reshape({1, :auto, 3, 3})
|> Grid.make_grid(grid_size: 16)
|> Nx.transpose(axes: [0, 2, 3, 1]),
randomMatrix_layer3
|> Nx.reshape({1, 120, :auto, 1}),
mat1_grad
|> Nx.reshape({1, 120, :auto, 1}),
randomMatrix_layer4 |> Nx.reshape({1, :auto, 120, 1}),
mat2_grad |> Nx.reshape({1, :auto, 120, 1})
]
|> KinoZoetrope.TensorStack.new(
size: [100, 100, 210, 210, 700, 700, 700, 700],
show_meta: false,
legend: false,
titel: "Gradients",
labels: [
"16 3x3 Kernels (L1)",
"Gradients (L1)",
"32*16 3x3 Kernels (L2)",
"Gradients (L2)",
"Matrix A (L3)",
"Matrix A Gradient (L3)",
"Matrix B (L4)",
"Matrix B Gradient (L4)"
]
)For actually performing the training step we now need to subtract all the gradient values from all the parameter values and then run all the calculations again, with the updated parameter values.
learning_rate = 0.002
val_n_grad =
Nx.Defn.jit(fn d,
l,
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
randomMatrix_layer3,
bias_layer3,
randomMatrix_layer4,
bias_layer4 ->
Nx.Defn.value_and_grad(
{
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
randomMatrix_layer3,
bias_layer3,
randomMatrix_layer4,
bias_layer4
},
fn {kernel16, b16, kernel32, b32, matrixA, m1b, matrixB, m2b} ->
Handwriting.mean_loss(
kernel16,
b16,
kernel32,
b32,
matrixA,
m1b,
matrixB,
m2b,
d,
l
)
end
)
end)
Stream.unfold(
# start with random weight values
{
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
randomMatrix_layer3,
bias_layer3,
randomMatrix_layer4,
bias_layer4
},
# iteration step
fn {layer1, bl1, layer2, bl2, layer3, bl3, layer4, bl4} ->
# calculate the loss value and all gradients
{loss,
{layer1_grad, bl1_grad, layer2_grad, bl2_grad, layer3_grad, bl3_grad, layer4_grad, bl4_grad}} =
val_n_grad.(
data,
labels_one_hot,
layer1,
bl1,
layer2,
bl2,
layer3,
bl3,
layer4,
bl4
)
# output calculated loss
{loss,
{
# subtract the gradient from the parameters (in small steps) for next step
layer1 |> Nx.subtract(layer1_grad |> Nx.multiply(learning_rate)),
bl1 |> Nx.subtract(bl1_grad |> Nx.multiply(learning_rate)),
layer2 |> Nx.subtract(layer2_grad |> Nx.multiply(learning_rate)),
bl2 |> Nx.subtract(bl2_grad |> Nx.multiply(learning_rate)),
layer3 |> Nx.subtract(layer3_grad |> Nx.multiply(learning_rate)),
bl3 |> Nx.subtract(bl3_grad |> Nx.multiply(learning_rate)),
layer4 |> Nx.subtract(layer4_grad |> Nx.multiply(learning_rate)),
bl4 |> Nx.subtract(bl4_grad |> Nx.multiply(learning_rate))
}}
end
)
|> Enum.take(10)
# [
# #Nx.Tensor<
# f32
# 13.78562
# >,
# #Nx.Tensor<
# f32
# 13.236868
# >,
# #Nx.Tensor<
# f32
# 12.755795
# >,
# #Nx.Tensor<
# f32
# 12.334583
# >,
# #Nx.Tensor<
# f32
# 11.855301
# >,
# #Nx.Tensor<
# f32
# 11.227427
# >,
# #Nx.Tensor<
# f32
# 10.616477
# >,
# #Nx.Tensor<
# f32
# 10.231294
# >,
# #Nx.Tensor<
# f32
# 9.860339
# >,
# #Nx.Tensor<
# f32
# 9.520728
# >
# ]Instead of running the training loop ourself we can load the parameters from precomputed files. There tensor values can just be stored and loaded as sequence of numbers. But we need to be careful to not mess up the order (row major vs column major)
dumped_files = [
{"dump_conv1.weight.txt", {16, 3, 3, 1}},
{"dump_conv1.bias.txt", {1, 4, 4, 1}},
{"dump_conv2.weight.txt", {32, 4 * 3, 4 * 3, 1}},
{"dump_conv2.bias.txt", {1, 4, 8, 1}},
{"dump_fc1.weight.txt", {1, 120, :auto, 1}},
{"dump_fc1.bias.txt", {1, 10, 12, 1}},
{"dump_fc2.weight.txt", {1, :auto, 100, 1}},
{"dump_fc2.bias.txt", {1, 1, 10, 1}}
]
pretrained =
for {d, s} <- dumped_files do
Kino.FS.file_path(d)
|> File.stream!()
|> CSV.decode!()
|> Stream.map(&Enum.map(&1, fn x -> String.to_float(x) end))
|> Enum.to_list()
|> Nx.tensor()
|> Nx.reshape(s)
end
pretrained
|> KinoZoetrope.TensorStack.new(
titel: "16 Random kernels",
cmap: :viridis,
labels:
for {d, _} <- dumped_files do
d
end
)defmodule HandwritingPredict do
import Nx.Defn
defn predict(
sixteenKernels,
sixteenBiases,
thirtyTwoKernels,
thirtyTwoBiases,
matrixA,
matABias,
matrixB,
matBBias,
data
) do
output =
data
|> Nx.rename([:batch, :channel, :x, :y])
|> Nx.pad(0, [{0, 0, 0}, {0, 0, 0}, {1, 1, 0}, {1, 1, 0}])
|> Nx.vectorize(:batch)
# |> Nx.conv(sixteenKernels, padding: [{1, 1}, {1, 1}])
|> NxSignal.Convolution.convolve(
sixteenKernels
|> Nx.vectorize(:channel),
mode: :valid
)
|> Nx.devectorize()
|> Nx.squeeze(axes: [2])
|> Nx.rename([:batch, :channel, :x, :y])
|> Nx.add(sixteenBiases)
|> Nx.max(0)
|> Nx.window_max({1, 1, 2, 2}, strides: [1, 1, 2, 2])
# |> Nx.conv(thirtyTwoKernels, padding: [{1, 1}, {1, 1}])
|> Nx.pad(0, [{0, 0, 0}, {0, 0, 0}, {1, 1, 0}, {1, 1, 0}])
|> Nx.vectorize(:batch)
|> NxSignal.Convolution.convolve(
thirtyTwoKernels
|> Nx.vectorize(:channel),
mode: :valid
)
|> Nx.devectorize()
|> Nx.squeeze(axes: [2])
|> Nx.rename([:batch, :channel, :x, :y])
|> Nx.add(thirtyTwoBiases)
|> Nx.max(0)
|> Nx.window_max({1, 1, 2, 2}, strides: [1, 1, 2, 2])
|> Nx.reshape({:auto, 7 * 7 * 32})
|> Nx.dot(matrixA)
|> Nx.add(matABias)
|> Nx.max(0)
|> Nx.dot(matrixB)
|> Nx.add(matBBias)
output_exp = Nx.exp(output)
output_exp
|> Nx.divide(Nx.sum(output_exp, axes: [1], keep_axes: true))
end
end[l1w, l1b, l2w, l2b, l3w, l3b, l4w, lwb] = pretrained
predictions =
HandwritingPredict.predict(
# l1w |> Nx.transpose(axes: [0,3,2,1]),
l1w
|> Nx.reshape({3, 3, 1, 16})
|> Nx.transpose(axes: [3, 2, 0, 1]),
l1b
|> Nx.reshape({1, 16, 1, 1}),
l2w
|> Nx.reshape({3, 3, 16, 32})
|> Nx.transpose(axes: [3, 2, 0, 1]),
l2b
|> Nx.reshape({1, 32, 1, 1}),
l3w
|> Nx.reshape({120, 1568})
|> Nx.transpose(axes: [1, 0]),
l3b
|> Nx.reshape({1, 120}),
l4w
|> Nx.reshape({10, 120})
|> Nx.transpose(axes: [1, 0]),
lwb |> Nx.reshape({10}),
data
)
all_loss =
Handwriting.loss(
# l1w |> Nx.transpose(axes: [0,3,2,1]),
l1w
|> Nx.reshape({3, 3, 1, 16})
|> Nx.transpose(axes: [3, 2, 0, 1])
|> Nx.reverse(axes: [-1, -2, -3]),
l1b
|> Nx.reshape({1, 16, 1, 1}),
l2w
|> Nx.reshape({3, 3, 16, 32})
|> Nx.transpose(axes: [3, 2, 0, 1])
|> Nx.reverse(axes: [-1, -2, -3]),
l2b
|> Nx.reshape({1, 32, 1, 1}),
l3w
|> Nx.reshape({120, 1568})
|> Nx.transpose(axes: [1, 0]),
l3b
|> Nx.reshape({1, 120}),
l4w
|> Nx.reshape({10, 120})
|> Nx.transpose(axes: [1, 0]),
lwb |> Nx.reshape({10}),
data,
labels_one_hot
)[
data |> Nx.transpose(axes: [0, 2, 3, 1]),
predictions |> Nx.reshape({1, 30, 10, 1}),
predictions
|> Nx.argmax(axis: 1)
|> Nx.reshape({1, 30, 1, 1})
|> Nx.transpose(axes: [0, 2, 1, 3]),
all_loss |> Nx.reshape({1, :auto, 10, 1}) |> Nx.log()
]
|> KinoZoetrope.TensorStack.new(
titel: "Random Matrix",
show_meta: false,
size: [300, 100, 400, 100],
labels: [
"Original Input",
"Prediction Distribution",
"Single Prediction",
"Loss"
],
legend: false,
markers: [
%{
attrs: %{
# SVG attributes
"stroke" => "magenta",
"stroke-width" => 2,
"fill" => "none",
"rx" => 10,
"ry" => 10
},
# indices of the images to add the markers to, here the first (0) and second (1) image
for: [2],
# list of {x,y} coordinates. The length if the list
# should match the number of frames of the image
points: 0..(Nx.size(labels) - 1) |> Enum.map(&{&1, 0})
},
%{
attrs: %{
# SVG attributes
"stroke" => "magenta",
"stroke-width" => 2,
"fill" => "none",
"rx" => 10,
"ry" => 10
},
# indices of the images to add the markers to, here the first (0) and second (1) image
for: [3],
# list of {x,y} coordinates. The length if the list
# should match the number of frames of the image
points: 0..(Nx.size(labels) - 1) |> Enum.map(&{rem(&1, 10), div(&1, 10)})
},
%{
attrs: %{
# SVG attributes
"stroke" => "magenta",
"stroke-width" => 2,
"fill" => "none",
"rx" => 10,
"ry" => 10
},
faded: true,
# indices of the images to add the markers to, here the first (0) and second (1) image
for: [1],
# list of {x,y} coordinates. The length if the list
# should match the number of frames of the image
points:
predictions
|> Nx.argmax(axis: 1)
|> Nx.to_flat_list()
|> Enum.with_index()
|> Enum.map(&{elem(&1, 0), elem(&1, 1)})
}
]
)