{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Importing your own neural net"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Verifying the example neural net was all well and good, but you probably want to verify your own neural net now. In this tutorial, we show you how to import the parameters for the feed-forward net in the introduction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "using MIPVerify\n",
    "# If you have an academic license, use `Gurobi` instead.\n",
    "using HiGHS\n",
    "using MAT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll download a `.mat` file containing the parameters of the sample neural net, containing two layers (exported from `tensorflow`). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dict{String, Any} with 20 entries:\n",
       "  \"fc1/weight\"           => Float32[-0.75335 -0.702841 … -0.0350578 -0.284728; …\n",
       "  \"logits/bias/Adam_1\"   => Float32[3.17918f-5 1.73752f-5 … 7.47935f-5 4.50664f…\n",
       "  \"logits/bias/Adam\"     => Float32[-0.000917174 0.000822852 … -0.00134091 -0.0…\n",
       "  \"fc1/bias\"             => Float32[0.763984 0.215094 … -0.702274 2.13576]\n",
       "  \"logits/weight/Adam_1\" => Float32[9.6672f-5 8.2316f-5 … 0.000832168 0.0001292…\n",
       "  \"fc2/weight/Adam_1\"    => Float32[3.93547f-6 3.1609f-5 … 7.98859f-7 6.52231f-…\n",
       "  \"fc2/bias\"             => Float32[0.657566 -1.76837 … 1.75472 0.168173]\n",
       "  \"beta1_power\"          => 0.0\n",
       "  \"logits/bias\"          => Float32[0.194171 -0.632471 … -1.13399 -1.00617]\n",
       "  \"fc2/bias/Adam\"        => Float32[-0.000178318 0.000325774 … -0.000285843 6.3…\n",
       "  \"fc1/weight/Adam\"      => Float32[0.0 0.0 … 4.9268f-6 -1.41152f-7; 5.39786f-1…\n",
       "  \"fc2/weight\"           => Float32[0.242356 0.800177 … 0.189305 -1.83328; -0.1…\n",
       "  \"logits/weight\"        => Float32[-0.0120639 0.0527637 … 0.458511 0.0677005; …\n",
       "  \"fc1/weight/Adam_1\"    => Float32[1.38988f-14 1.15851f-15 … 9.08535f-9 5.8825…\n",
       "  \"fc2/weight/Adam\"      => Float32[0.000683748 -0.000629174 … 0.000171999 -7.7…\n",
       "  \"beta2_power\"          => 0.0\n",
       "  \"fc1/bias/Adam\"        => Float32[0.000136568 7.1624f-5 … 0.000106514 0.00081…\n",
       "  \"logits/weight/Adam\"   => Float32[-0.00237726 0.000125507 … -0.000695765 -0.0…\n",
       "  \"fc1/bias/Adam_1\"      => Float32[1.77318f-6 2.53318f-6 … 3.44693f-6 2.25489f…\n",
       "  \"fc2/bias/Adam_1\"      => Float32[7.38528f-6 5.92213f-6 … 1.58686f-6 4.14633f…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "param_dict = Base.download(\"https://github.com/vtjeng/MIPVerify_data/raw/master/weights/mnist/n1.mat\") |> matread"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layer 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's begin by importing the parameters for the first fully connected layer, which has 784 inputs (corresponding to a flattened 28x28 image) and 40 outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic Approach\n",
    "\n",
    "We begin with a basic approach where we extract the weights and the biases of the fully connected layer seperately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "784×40 Matrix{Float32}:\n",
       " -0.75335   -0.702841   0.00934679   0.0424287   …  -0.0350578   -0.284728\n",
       " -1.26091   -0.975876  -0.0690286   -0.145539       -0.139512    -0.494751\n",
       " -1.30244   -1.54214   -0.063707    -0.306707       -0.0144018   -0.644423\n",
       " -1.07623   -0.899058  -0.0936826   -0.025368        0.0716064   -0.357508\n",
       " -1.09825   -1.49246   -0.137173    -0.272637       -0.211558    -0.73989\n",
       " -0.96502   -0.886079  -0.081703    -0.0568857   …   0.00850921  -0.274247\n",
       " -0.733129  -0.931515  -0.0551823    0.0135264       0.138682    -0.322842\n",
       " -0.991218  -0.86704   -0.036489    -0.0409108       0.0330326   -0.472234\n",
       " -1.09089   -1.31528   -0.0488595    0.00274563     -0.160661    -0.603701\n",
       " -1.0728    -1.15951   -0.0882954   -0.10305        -0.138768    -0.700878\n",
       "  ⋮                                              ⋱               \n",
       " -1.42311   -1.20815   -0.0115304    0.111935    …  -0.41617     -0.518008\n",
       " -1.9324    -1.53047   -0.109568    -0.0240137      -0.608832    -1.12603\n",
       " -1.51948   -1.42803    0.248483    -0.0787689      -0.0857386   -0.729084\n",
       " -1.63084   -1.27073    0.158391    -0.0356018      -0.323439    -0.692301\n",
       " -1.01737   -1.03474   -0.0673599    0.0771562      -0.133131    -0.396228\n",
       " -0.80131   -0.675911  -0.0228199    0.0430513   …  -0.0688749   -0.451306\n",
       " -0.943113  -0.780626  -0.0223145    0.0169085       0.237589    -0.340899\n",
       " -0.885898  -0.713929  -0.0602917   -0.021813       -0.111022    -0.320723\n",
       " -0.963221  -0.951116  -0.0672328   -0.0266708       0.137328    -0.312931"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fc1_weight = param_dict[\"fc1/weight\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×40 Matrix{Float32}:\n",
       " 0.763984  0.215094  4.37897  -0.164343  …  4.89644  -0.702274  2.13576"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fc1_bias = param_dict[\"fc1/bias\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We group the weights and biases in a `Linear`.\n",
    "\n",
    "_(NB: We have to flatten the bias layer using `dropdims` since `Linear` expects a 1-D array for the bias.)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(784 -> 40)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fc1_manual = Linear(fc1_weight, dropdims(fc1_bias, dims=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That was a lot to remember. Wouldn't it be nice if there was a helper function to take care of all that?\n",
    "\n",
    "### With Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(784 -> 40)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fc1 = get_matrix_params(param_dict, \"fc1\", (784, 40))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`get_matrix_params` requires that 1) you specify the expected size of the layer, and 2) your weight and bias arrays following the naming convention outlined in the [documentation](https://vtjeng.com/MIPVerify.jl/stable/utils/import_weights/#MIPVerify.get_matrix_params-Tuple{Dict{String},%20String,%20Tuple{Int64,%20Int64}})."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a quick check, you can verify that the parameters we get from both methods are equal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fc1.matrix == fc1_manual.matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fc1.bias == fc1_manual.bias"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing the rest of the layers\n",
    "\n",
    "Since we followed the naming convention required by `get_matrix_params` when exporting our neural net parameters as a `.mat` file, importing the rest of the neural net is relatively straightforward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(40 -> 20)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fc2 = get_matrix_params(param_dict, \"fc2\", (40, 20))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(20 -> 10)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "logits = get_matrix_params(param_dict, \"logits\", (20, 10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Composing the network\n",
    "\n",
    "We now put the entire network together. We need to start by flattening the input since the input images are provided as a 4-dimensional tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "sequential net MNIST.n1\n",
       "  (1) Flatten(): flattens 4 dimensional input, with dimensions permuted according to the order [4, 3, 2, 1]\n",
       "  (2) Linear(784 -> 40)\n",
       "  (3) ReLU()\n",
       "  (4) Linear(40 -> 20)\n",
       "  (5) ReLU()\n",
       "  (6) Linear(20 -> 10)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n1 = Sequential([\n",
    "        Flatten(4),\n",
    "        fc1,\n",
    "        # you can always use interval arithmetic for the first layer\n",
    "        ReLU(interval_arithmetic),\n",
    "        fc2,\n",
    "        ReLU(),\n",
    "        logits\n",
    "    ], \"MNIST.n1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There we go! Now it's your turn to try to verify your own neural network."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.7.1",
   "language": "julia",
   "name": "julia-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}