{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "149938de",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "773c1396",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, test_data = tf.keras.datasets.mnist.load_data(path=\"mnist.npz\")\n",
    "x_train = train_data[0].reshape(-1,784).astype(float) / 255\n",
    "y_train = train_data[1]\n",
    "print(y_train[0])\n",
    "plt.imshow(train_data[0][0], cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b521b78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# source: https://keras.io/guides/training_with_built_in_methods/\n",
    "\n",
    "inputs = keras.Input(shape=(784,), name=\"digits\")\n",
    "x = layers.Dense(32, activation=\"relu\", name=\"dense_1\")(inputs)\n",
    "x = layers.Dense(32, activation=\"relu\", name=\"dense_2\")(x)\n",
    "outputs = layers.Dense(10, activation=\"softmax\", name=\"predictions\")(x)\n",
    "model = keras.Model(inputs=inputs, outputs=outputs, name=\"mnist_model\")\n",
    "loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "#def loss(true, pred):\n",
    "#    val = 0.\n",
    "#    for i in range(len(true)):\n",
    "#        val -= np.log(pred[i, true[i]])\n",
    "#    return val / len(true)\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b18b6bdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# untrained model\n",
    "logits = model(test_data[0], training=True)\n",
    "np.argmax(logits, 1)[0:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d74f67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# what we would like to get\n",
    "test_data[1][0:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c0073e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Computing a stochastic gradient\n",
    "\n",
    "# Source: https://keras.io/guides/writing_a_training_loop_from_scratch/\n",
    "\n",
    "# Example of gradient computation using automatic differentiation\n",
    "\n",
    "x_train_batch = x_train[:8]\n",
    "y_train_batch = y_train[:8]\n",
    "with tf.GradientTape() as tape:\n",
    "        # Run the forward pass of the layer.\n",
    "        # The operations that the layer applies\n",
    "        # to its inputs are going to be recorded\n",
    "        # on the GradientTape.\n",
    "        logits = model(x_train_batch)  # Logits for this minibatch\n",
    "\n",
    "        # Compute the loss value for this minibatch.\n",
    "        loss_value = loss(y_train_batch, logits)\n",
    "\n",
    "# Use the gradient tape to automatically retrieve\n",
    "# the gradients of the trainable variables with respect to the loss.\n",
    "grad = tape.gradient(loss_value, model.trainable_weights)\n",
    "\n",
    "## Updating weights\n",
    "\n",
    "def update_weights(model, update):\n",
    "    new_weights = model.trainable_weights.copy()\n",
    "    for i in range(len(new_weights)):\n",
    "        new_weights[i] = new_weights[i] + update[i]\n",
    "    model.set_weights(new_weights)\n",
    "\n",
    "\n",
    "## Evaluate model\n",
    "samples_for_evaluation = np.random.randint(0, 10000, 1000)\n",
    "logits = model(x_train[samples_for_evaluation])\n",
    "loss(y_train[samples_for_evaluation], logits).numpy()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5a9ead7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Stochastic gradient algorithm\n",
    "\n",
    "# TO IMPLEMENT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16641e62",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec79dd02",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ba54737",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3845cef",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d9bf5cc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24d39aac",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}