Solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b08e802e-b75c-421d-a9c8-53ddfb23c849",
   "metadata": {},
   "source": [
    "# AWID RNN\n",
    "#### Tomasz Mycielski"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18edbd53-bef6-49f7-bfb7-fe514fc5c7e3",
   "metadata": {},
   "source": [
    "## External packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0b180ec2-a9b3-4e67-9b40-199ba9211308",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m   Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.10/Project.toml`\n",
      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.10/Manifest.toml`\n",
      "\u001b[32m\u001b[1m   Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.10/Project.toml`\n",
      "\u001b[32m\u001b[1m  No Changes\u001b[22m\u001b[39m to `~/.julia/environments/v1.10/Manifest.toml`\n"
     ]
    }
   ],
   "source": [
    "using Pkg\n",
    "Pkg.add(\"MLDatasets\")\n",
    "Pkg.add(\"Flux\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afe37d4e-8374-4251-b102-8e2cf1f42b26",
   "metadata": {},
   "source": [
    "---\n",
    "## Boilerplate\n",
    "(from lecture #4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "753e9d65-07a3-4e89-8197-3187b378a08d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import Base: *\n",
    "import LinearAlgebra: mul!, diagm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e34b9938-25d8-4da2-a2ab-8da41a69fc79",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "backward! (generic function with 4 methods)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "abstract type GraphNode end\n",
    "abstract type Operator <: GraphNode end\n",
    "\n",
    "struct Constant{T} <: GraphNode\n",
    "    output :: T\n",
    "end\n",
    "\n",
    "mutable struct Variable <: GraphNode\n",
    "    output :: Any\n",
    "    gradient :: Any\n",
    "    name :: String\n",
    "    Variable(output; name=\"?\") = new(output, nothing, name)\n",
    "end\n",
    "\n",
    "mutable struct ScalarOperator{F} <: Operator\n",
    "    inputs :: Any\n",
    "    output :: Any\n",
    "    gradient :: Any\n",
    "    name :: String\n",
    "    ScalarOperator(fun, inputs...; name=\"?\") = new{typeof(fun)}(inputs, nothing, nothing, name)\n",
    "end\n",
    "\n",
    "mutable struct BroadcastedOperator{F} <: Operator\n",
    "    inputs :: Any\n",
    "    output :: Any\n",
    "    gradient :: Any\n",
    "    name :: String\n",
    "    BroadcastedOperator(fun, inputs...; name=\"?\") = new{typeof(fun)}(inputs, nothing, nothing, name)\n",
    "end\n",
    "\n",
    "import Base: show, summary\n",
    "show(io::IO, x::ScalarOperator{F}) where {F} = print(io, \"op \", x.name, \"(\", F, \")\");\n",
    "show(io::IO, x::BroadcastedOperator{F}) where {F} = print(io, \"op.\", x.name, \"(\", F, \")\");\n",
    "show(io::IO, x::Constant) = print(io, \"const \", x.output)\n",
    "show(io::IO, x::Variable) = begin\n",
    "    print(io, \"var \", x.name);\n",
    "    print(io, \"\\n ┣━ ^ \"); summary(io, x.output)\n",
    "    print(io, \"\\n ┗━ ∇ \");  summary(io, x.gradient)\n",
    "end\n",
    "\n",
    "function visit(node::GraphNode, visited, order)\n",
    "    if node ∈ visited\n",
    "    else\n",
    "        push!(visited, node)\n",
    "        push!(order, node)\n",
    "    end\n",
    "    return nothing\n",
    "end\n",
    "    \n",
    "function visit(node::Operator, visited, order)\n",
    "    if node ∈ visited\n",
    "    else\n",
    "        push!(visited, node)\n",
    "        for input in node.inputs\n",
    "            visit(input, visited, order)\n",
    "        end\n",
    "        push!(order, node)\n",
    "    end\n",
    "    return nothing\n",
    "end\n",
    "\n",
    "function topological_sort(head::GraphNode)\n",
    "    visited = Set()\n",
    "    order = Vector()\n",
    "    visit(head, visited, order)\n",
    "    return order\n",
    "end\n",
    "\n",
    "reset!(node::Constant) = nothing\n",
    "reset!(node::Variable) = node.gradient = nothing\n",
    "reset!(node::Operator) = node.gradient = nothing\n",
    "\n",
    "compute!(node::Constant) = nothing\n",
    "compute!(node::Variable) = nothing\n",
    "compute!(node::Operator) = let\n",
    "    node.output = forward(node, [input.output for input in node.inputs]...)\n",
    "end\n",
    "\n",
    "function forward!(order::Vector)\n",
    "    for node in order\n",
    "        compute!(node)\n",
    "        reset!(node)\n",
    "    end\n",
    "    return last(order).output\n",
    "end\n",
    "\n",
    "update!(node::Constant, gradient) = nothing\n",
    "update!(node::GraphNode, gradient) = if isnothing(node.gradient)\n",
    "    node.gradient = gradient else node.gradient .+= gradient\n",
    "end\n",
    "\n",
    "function backward!(order::Vector; seed=1.0)\n",
    "    result = last(order)\n",
    "    result.gradient = seed\n",
    "    #@assert length(result.output) == 1 \"Gradient is defined only for scalar functions\"\n",
    "    for node in reverse(order)\n",
    "        backward!(node)\n",
    "    end\n",
    "    return nothing\n",
    "end\n",
    "\n",
    "function backward!(node::Constant) end\n",
    "function backward!(node::Variable) end\n",
    "function backward!(node::Operator)\n",
    "    inputs = node.inputs\n",
    "    gradients = backward(node, [input.output for input in inputs]..., node.gradient)\n",
    "    for (input, gradient) in zip(inputs, gradients)\n",
    "        update!(input, gradient)\n",
    "    end\n",
    "    return nothing\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad33888a-d8f3-4918-a63b-47cd477cd6a7",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3b4e8ca4-a816-488b-b077-460a642786c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10000-element DataLoader(::Tuple{Matrix{Float64}, OneHotArrays.OneHotMatrix{UInt32, Vector{UInt32}}}, shuffle=true)\n",
       "  with first element:\n",
       "  (784×1 Matrix{Float64}, 10×1 OneHotMatrix(::Vector{UInt32}) with eltype Bool,)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import MLDatasets\n",
    "train_data = MLDatasets.MNIST(split=:train)\n",
    "test_data  = MLDatasets.MNIST(split=:test)\n",
    "\n",
    "import Flux\n",
    "function loader(data; batchsize::Int=1)\n",
    "    x1dim = reshape(data.features, 28 * 28, :) |> Matrix{Float64}# reshape 28×28 pixels into a vector of pixels\n",
    "    yhot  = Flux.onehotbatch(data.targets, 0:9) # make a 10×60000 OneHotMatrix\n",
    "    Flux.DataLoader((x1dim, yhot); batchsize, shuffle=true)\n",
    "end\n",
    "\n",
    "train_loader = loader(train_data)\n",
    "test_loader = loader(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f695070-cdc8-439d-ae99-bf00b2747e35",
   "metadata": {},
   "source": [
    "---\n",
    "## Implemented operations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe33d3ce-2684-4aa9-b865-cee97db2cafe",
   "metadata": {},
   "source": [
    "### Multiplication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "72b086d1-a1c2-4dbf-8ee7-de1c275cd525",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "backward (generic function with 2 methods)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "*(A::GraphNode, x::GraphNode) = BroadcastedOperator(mul!, A, x)\n",
    "forward(::BroadcastedOperator{typeof(mul!)}, A, x) = return A * x\n",
    "backward(::BroadcastedOperator{typeof(mul!)}, A, x, g) = tuple(g * x', A' * g)\n",
    "\n",
    "Base.Broadcast.broadcasted(*, x::GraphNode, y::GraphNode) = BroadcastedOperator(*, x, y)\n",
    "forward(::BroadcastedOperator{typeof(*)}, x, y) = return x .* y\n",
    "backward(node::BroadcastedOperator{typeof(*)}, x, y, g) = let\n",
    "    𝟏 = ones(length(node.output))\n",
    "    Jx = diagm(y .* 𝟏)\n",
    "    Jy = diagm(vec(x .* 𝟏))\n",
    "    tuple(Jx' * g, Jy' * g)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "233618e1-c1ff-4aa9-a89b-059bc710d666",
   "metadata": {},
   "source": [
    "### Addition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "78871b70-cd2e-428d-83a9-5fc1ebaaeb9e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "backward (generic function with 3 methods)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Base.Broadcast.broadcasted(+, x::GraphNode, y::GraphNode) = BroadcastedOperator(+, x, y)\n",
    "forward(::BroadcastedOperator{typeof(+)}, x, y) = return x .+ y\n",
    "backward(::BroadcastedOperator{typeof(+)}, x, y, g) = tuple(g, g)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2454746-5684-4e69-bc61-58149ad31061",
   "metadata": {},
   "source": [
    "### Summation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5a0a83d2-ff92-432b-8fc5-af3099159fde",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "backward (generic function with 4 methods)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import Base: sum\n",
    "sum(x::GraphNode) = BroadcastedOperator(sum, x)\n",
    "forward(::BroadcastedOperator{typeof(sum)}, x) = return sum(x)\n",
    "backward(::BroadcastedOperator{typeof(sum)}, x, g) = let\n",
    "    𝟏 = ones(length(x))\n",
    "    J = 𝟏'\n",
    "    tuple(J' * g)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f39d577-0e33-4d72-a9aa-f76e500294b3",
   "metadata": {},
   "source": [
    "### Tanh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b6d4dc51-7385-485e-aa4e-c8a6c117d9a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import Base: tanh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ced4124f-f44a-450a-b9a0-9b996b5abdb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "backward (generic function with 5 methods)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tanh(x::GraphNode) = BroadcastedOperator(tanh, x)\n",
    "forward(::BroadcastedOperator{typeof(tanh)}, x) = return tanh.(x)\n",
    "backward(node::BroadcastedOperator{typeof(tanh)}, x, g) = return tuple((1 .- tanh.(x).^2) .* g)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28a04a44-6c25-4675-bd7b-036eb139bbc3",
   "metadata": {},
   "source": [
    "### Softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f2074557-bd0a-4f80-9e1e-8f4811976903",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "backward (generic function with 6 methods)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Softmax(x::GraphNode) = BroadcastedOperator(Softmax, x)\n",
    "forward(::BroadcastedOperator{typeof(Softmax)}, x) = return exp.(x) ./ sum(exp.(x))\n",
    "backward(node::BroadcastedOperator{typeof(Softmax)}, x, g) = let\n",
    "    y = node.output\n",
    "    J = diagm(y) .- y * y'\n",
    "    tuple(J' * g)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed463431-e5af-4610-adf8-3d9fe5534952",
   "metadata": {},
   "source": [
    "### Log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e5691a19-641b-41f5-ba98-1f168f8bce22",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "backward (generic function with 7 methods)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Base.Broadcast.broadcasted(log, x::GraphNode) = BroadcastedOperator(log, x)\n",
    "forward(::BroadcastedOperator{typeof(log)}, x) = return log.(x)\n",
    "backward(::BroadcastedOperator{typeof(log)}, x, g) = let\n",
    "    tuple(g ./ x)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f71945a3-1a97-4f78-91e9-caca4c47f332",
   "metadata": {},
   "source": [
    "---\n",
    "## Network"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41b197da-3786-422c-85d4-909fc83b9487",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a92206b0-1219-4a9e-91e8-7b3d669f9b31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "INPUT_SIZE = 196\n",
    "HIDDEN_SIZE = 64\n",
    "OUTPUT_SIZE = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7d09380-cf02-46e2-b1db-912cb2f89eff",
   "metadata": {},
   "source": [
    "### Weights init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "faa1f90d-66e0-43e3-aa17-bdbb61457168",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "glorot (generic function with 1 method)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function glorot(size_a, size_b)\n",
    "    total_size = size_a + size_b\n",
    "    denum = sqrt(total_size)\n",
    "    return randn(size_a, size_b) .* ( sqrt(6)/denum )\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2d3e6a79-5cef-4524-91ad-0b04d7ade0dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "var ?\n",
       " ┣━ ^ 10-element Vector{Float32}\n",
       " ┗━ ∇ Nothing"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Wi = Variable(glorot(HIDDEN_SIZE, INPUT_SIZE))\n",
    "Wh = Variable(glorot(HIDDEN_SIZE, HIDDEN_SIZE))\n",
    "Wo = Variable(glorot(OUTPUT_SIZE, HIDDEN_SIZE))\n",
    "\n",
    "Bh = Variable(zeros(HIDDEN_SIZE) |> Vector{Float32})\n",
    "Bo = Variable(zeros(OUTPUT_SIZE) |> Vector{Float32})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b11d7fa-e597-4ff5-ba57-ab3393a359b3",
   "metadata": {},
   "source": [
    "### Loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a9e4a079-fb9f-432e-b49a-2e479e19e190",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "cross_entropy_loss (generic function with 1 method)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function cross_entropy_loss(prediction, label)\n",
    "    return Constant(-1) .* sum(Variable(label) .* log.(prediction))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0752f748-6852-459c-927f-612619c13ac6",
   "metadata": {},
   "source": [
    "### Network structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "37b2b64b-3e4d-43c7-9ac9-e1d98f761ed0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "net (generic function with 1 method)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function net(sample, input_weights, hidden_weights, output_weights, label)\n",
    "    step_input = Variable(view(sample, 1:196))\n",
    "\n",
    "    hidden_state = tanh(Wi * step_input .+ Bh)\n",
    "\n",
    "    step_input = Variable(view(sample, 197:392))\n",
    "    hidden_state = tanh(Wi * step_input .+ Wh * hidden_state .+ Bh)\n",
    "\n",
    "    step_input = Variable(view(sample, 393:588))\n",
    "    hidden_state = tanh(Wi * step_input .+ Wh * hidden_state .+ Bh)\n",
    "\n",
    "    step_input = Variable(view(sample, 589:784))\n",
    "    hidden_state = tanh(Wi * step_input .+ Wh * hidden_state .+ Bh)\n",
    "\n",
    "    prediction = Softmax(Wo * hidden_state .+ Bo)\n",
    "\n",
    "    E = cross_entropy_loss(prediction, label)\n",
    "\n",
    "    return topological_sort(E), prediction\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0aa35f0b-ac46-411a-b148-660b8caa4554",
   "metadata": {},
   "source": [
    "---\n",
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "430a45cd-56f0-4c68-8d99-cad0654e70eb",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4a8078a2-fab1-4585-804b-6ee58bc6710f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "STEP_SIZE = Float64(15e-3)\n",
    "EPOCHS = 5\n",
    "BATCH_SIZE = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bb54209-0287-417d-885d-b3f8d0505ecb",
   "metadata": {},
   "source": [
    "### Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "35d9532f-561e-416d-8f5c-8ab8c4664d5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 17.366283 seconds (75.11 M allocations: 33.080 GiB, 9.63% gc time, 19.84% compilation time)\n",
      " 14.805266 seconds (61.64 M allocations: 32.206 GiB, 12.93% gc time)\n",
      " 13.605452 seconds (61.64 M allocations: 32.206 GiB, 9.97% gc time)\n",
      " 13.110781 seconds (61.64 M allocations: 32.206 GiB, 8.76% gc time)\n",
      " 13.177806 seconds (61.64 M allocations: 32.206 GiB, 9.32% gc time)\n"
     ]
    }
   ],
   "source": [
    "for epoch_index in range(start=1, stop=EPOCHS)\n",
    "    Wi_grad_agg = zeros(HIDDEN_SIZE, INPUT_SIZE)\n",
    "    Wh_grad_agg = zeros(HIDDEN_SIZE, HIDDEN_SIZE)\n",
    "    Wo_grad_agg = zeros(OUTPUT_SIZE, HIDDEN_SIZE)\n",
    "\n",
    "    Bh_grad_agg = zeros(HIDDEN_SIZE)\n",
    "    Bo_grad_agg = zeros(OUTPUT_SIZE)\n",
    "    \n",
    "    @time for (index, (s, l)) in enumerate(train_loader)\n",
    "        graph, predicted = net(s, Wi, Wh, Wo, l)\n",
    "        forward!(graph)\n",
    "        backward!(graph)\n",
    "    \n",
    "        Wi_grad_agg .+= Wi.gradient\n",
    "        Wh_grad_agg .+= Wh.gradient\n",
    "        Wo_grad_agg .+= Wo.gradient\n",
    "\n",
    "        Bh_grad_agg .+= Bh.gradient\n",
    "        Bo_grad_agg .+= Bo.gradient\n",
    "        \n",
    "        if index % BATCH_SIZE == 0\n",
    "            Wi.output .-= ((Wi_grad_agg/BATCH_SIZE) * STEP_SIZE)\n",
    "            Wh.output .-= ((Wh_grad_agg/BATCH_SIZE) * STEP_SIZE)\n",
    "            Wo.output .-= ((Wo_grad_agg/BATCH_SIZE) * STEP_SIZE)\n",
    "            Wi_grad_agg .*= 0\n",
    "            Wh_grad_agg .*= 0\n",
    "            Wo_grad_agg .*= 0\n",
    "\n",
    "            Bh.output .-= ((Bh_grad_agg/BATCH_SIZE) * STEP_SIZE)\n",
    "            Bo.output .-= ((Bo_grad_agg/BATCH_SIZE) * STEP_SIZE)\n",
    "            Bh_grad_agg .*= 0\n",
    "            Bo_grad_agg .*= 0\n",
    "\n",
    "        end\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "402c522c-09e3-431f-a02b-5c28bc73c1cc",
   "metadata": {},
   "source": [
    "- 13 seconds per epoch\n",
    "- 61.64 milion allocations (32.2 GiB) per epoch\n",
    "- 8% gc time\n",
    "\n",
    "Total training time is about a minute"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "980ffa82-f55e-4640-ae43-80aab20b0e47",
   "metadata": {},
   "source": [
    "### Testing loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fd8b3d40-e4c9-41ed-92a9-f48904d2f26e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  2.248723 seconds (10.37 M allocations: 5.349 GiB, 10.20% gc time, 2.10% compilation time)\n",
      "\n",
      "Test accuracy:\n",
      "92.16%\n"
     ]
    }
   ],
   "source": [
    "correct = 0\n",
    "total = 0\n",
    "\n",
    "@time for (index, (s, l)) in enumerate(test_loader)\n",
    "    graph, predicted = net(s, Wi, Wh, Wo, l)\n",
    "    forward!(graph)\n",
    "    backward!(graph)\n",
    "    if argmax(predicted.output) == argmax(l)[1]\n",
    "        correct += 1\n",
    "    end\n",
    "    total += 1\n",
    "end\n",
    "\n",
    "println(\"\\nTest accuracy:\")\n",
    "println(round(correct/total*100, digits=4), \"%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebfcbacb-73df-44fe-bda6-9a55a93f3414",
   "metadata": {},
   "source": [
    "**Decent**, but I've seen it get as high as $96.5\\%$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4cd4f72-d37a-4cae-9b34-eeedebee6a4d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.3",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}