# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] colab_type="text" id="Yu7u1mNfnxVP"
# **Copyright 2019 The Sonnet Authors. All Rights Reserved.**
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# ---

# + [markdown] colab_type="text" id="WAfR3cvnoGMB"
# # Preamble

# + colab={} colab_type="code" id="4FqOAJb_jJR9"
import sys
assert sys.version_info >= (3, 6), "Sonnet 2 requires Python >=3.6"

# + colab={"base_uri": "https://localhost:8080/", "height": 136} colab_type="code" id="XnWX2azUDuCl" outputId="8516864f-06cb-4dd5-946c-adeae1e17f3a"
# !pip install dm-sonnet tqdm

# + colab={} colab_type="code" id="mn5ofK4-D1Qk"
import sonnet as snt
import tensorflow as tf
import tensorflow_datasets as tfds

# + colab={"base_uri": "https://localhost:8080/", "height": 51} colab_type="code" id="Rpp_houJEHr9" outputId="9e7597ba-d8a5-483e-b415-5729ef3102e8"
print("TensorFlow version: {}".format(tf.__version__))
print("    Sonnet version: {}".format(snt.__version__))

# + [markdown] colab_type="text" id="5RmHUmz1padR"
# Finally lets take a quick look at the GPUs we have available:

# + colab={"base_uri": "https://localhost:8080/", "height": 34} colab_type="code" id="TXoxEvKepdw2" outputId="da76b78b-274e-462d-fa47-89d626a4370f"
# !grep Model: /proc/driver/nvidia/gpus/*/information | awk '{$1="";print$0}'

# + [markdown] colab_type="text" id="UYYmqvOKfNbk"
# # Dataset
#
# We need to get our dataset in a state where we can iterate over it easily. The TensorFlow Datasets package provides a simple API for this. It will download the dataset and prepare it for us to speedily process on a GPU. We can also add our own pre-processing functions to mutate the dataset before our model sees it:

# + colab={} colab_type="code" id="UkBRriaQEr4z"
batch_size = 100

def process_batch(images, labels):
  images = tf.squeeze(images, axis=[-1])
  images = tf.cast(images, dtype=tf.float32)
  images = ((images / 255.) - .5) * 2.
  return images, labels

def mnist(split):
  dataset = tfds.load("mnist", split=split, as_supervised=True)
  dataset = dataset.map(process_batch)
  dataset = dataset.batch(batch_size)
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  dataset = dataset.cache()
  return dataset

mnist_train = mnist("train").shuffle(10)
mnist_test = mnist("test")

# + [markdown] colab_type="text" id="JfOCWVGEfgcq"
# MNIST contains `28x28` greyscale handwritten digits. Let's take a look at one:

# + colab={"base_uri": "https://localhost:8080/", "height": 269} colab_type="code" id="I_yM0TVjFCZq" outputId="54959260-c8fc-4e39-ea54-e566e9122893"
import matplotlib.pyplot as plt

images, _ = next(iter(mnist_test))
plt.imshow(images[0]);


# + [markdown] colab_type="text" id="d7bsizs5gK3K"
# # Sonnet
#
# The next step is to define a model. In Sonnet everything that contains TensorFlow variables (`tf.Variable`) extends `snt.Module`, this includes low level neural network components (e.g. `snt.Linear`,  `snt.Conv2D`), larger nets containing subcomponents (e.g. `snt.nets.MLP`), optimizers (e.g. `snt.optimizers.Adam`) and whatever else you can think of.
#
# Modules provide a simple abstraction for storing parameters (and `Variable`s used for other purposes, like for storing moving avergages in `BatchNorm`).
#
# To find all the parameters for a given module, simply do: `module.variables`. This will return a `tuple` of all the parameters that exist for this module, or any module it references:

# + [markdown] colab_type="text" id="GrN37pi1o4HT"
# ## Building the model

# + [markdown] colab_type="text" id="c6XoN56S2lSW"
# In Sonnet you build neural networks out of `snt.Module`s. In this case we'll build a multi-layer perceptron as a new class with a `__call__` method that computes the logits by passing the input through a number of fully connected layers, with a ReLU non-linearity.

# + colab={} colab_type="code" id="hgjyB9yhFclD"
class MLP(snt.Module):

  def __init__(self):
    super(MLP, self).__init__()
    self.flatten = snt.Flatten()
    self.hidden1 = snt.Linear(1024, name="hidden1")
    self.hidden2 = snt.Linear(1024, name="hidden2")
    self.logits = snt.Linear(10, name="logits")

  def __call__(self, images):
    output = self.flatten(images)
    output = tf.nn.relu(self.hidden1(output))
    output = tf.nn.relu(self.hidden2(output))
    output = self.logits(output)
    return output


# + [markdown] colab_type="text" id="0i03px8y8gf7"
# Now we'll create an instance of our class whose weights will be randomly initialized. We'll train this MLP such that it learns to recognize digits in the MNIST dataset.

# + colab={"base_uri": "https://localhost:8080/", "height": 34} colab_type="code" id="XqL8oIMqGAnU" outputId="96efc445-d7bb-45db-f6b1-5aa4a8b72dcb"
mlp = MLP()
mlp

# + [markdown] colab_type="text" id="snzkUUh9oXPy"
# ## Using the model

# + [markdown] colab_type="text" id="On8wI6VwpDPm"
# Let's feed an example input through the model and see what it predicts. Since the model is randomly initialized there is a 1/10 chance that it will predict the right class!

# + colab={"base_uri": "https://localhost:8080/", "height": 286} colab_type="code" id="4T-qmIc0GHfP" outputId="316f41f2-319c-452e-a132-5006b732e86b"
images, labels = next(iter(mnist_test))
logits = mlp(images)
  
prediction = tf.argmax(logits[0]).numpy()
actual = labels[0].numpy()
print("Predicted class: {} actual class: {}".format(prediction, actual))
plt.imshow(images[0]);

# + [markdown] colab_type="text" id="V297xpzfobXK"
# ## Training the model

# + [markdown] colab_type="text" id="WTrv-jn4pPSx"
# To train the model we need an optimizer. For this simple example we'll use Stochastic Gradient Descent which is implemented in the `SGD` optimizer. To compute gradients we'll use a `tf.GradientTape` which allows us to selectively record gradients only for the computation we want to back propagate through:

# + cellView="form" colab={} colab_type="code" id="V7gi8NQ-WZOl"
#@title Utility function to show progress bar.
from tqdm import tqdm

# MNIST training set has 60k images.
num_images = 60000

def progress_bar(generator):
  return tqdm(
      generator,
      unit='images',
      unit_scale=batch_size,
      total=(num_images // batch_size) * num_epochs)



# + colab={"base_uri": "https://localhost:8080/", "height": 85} colab_type="code" id="UUkshshiK6Eq" outputId="edf09c6d-cc08-482d-98d8-45778082ed0b"
opt = snt.optimizers.SGD(learning_rate=0.1)

num_epochs = 10

def step(images, labels):
  """Performs one optimizer step on a single mini-batch."""
  with tf.GradientTape() as tape:
    logits = mlp(images)
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
                                                          labels=labels)
    loss = tf.reduce_mean(loss)

  params = mlp.trainable_variables
  grads = tape.gradient(loss, params)
  opt.apply(grads, params)
  return loss

for images, labels in progress_bar(mnist_train.repeat(num_epochs)):
  loss = step(images, labels)

print("\n\nFinal loss: {}".format(loss.numpy()))

# + [markdown] colab_type="text" id="2K0_eoR8og-G"
# ## Evaluating the model

# + [markdown] colab_type="text" id="Cm_9RMJopgWc"
# We'll do very simple analysis of the model to get a feeling for how well it does against this dataset:

# + colab={"base_uri": "https://localhost:8080/", "height": 34} colab_type="code" id="PM7IPcOeXtxH" outputId="fc983d03-f92c-4002-f7df-514c251bd300"
total = 0
correct = 0
for images, labels in mnist_test:
  predictions = tf.argmax(mlp(images), axis=1)
  correct += tf.math.count_nonzero(tf.equal(predictions, labels))
  total += images.shape[0]

print("Got %d/%d (%.02f%%) correct" % (correct, total, correct / total * 100.))


# + [markdown] colab_type="text" id="Lnkc55PtqA_I"
# To understand the result a bit better, lets take a look at a small sample of where the model correctly identified the digits:

# + cellView="form" colab={} colab_type="code" id="eFro_RB4YR-X"
#@title Utility function to show a sample of images.
def sample(correct, rows, cols):
  n = 0

  f, ax = plt.subplots(rows, cols)
  if rows > 1:    
    ax = tf.nest.flatten([tuple(ax[i]) for i in range(rows)])
  f.set_figwidth(14)
  f.set_figheight(4 * rows)


  for images, labels in mnist_test:
    predictions = tf.argmax(mlp(images), axis=1)
    eq = tf.equal(predictions, labels)
    for i, x in enumerate(eq):
      if x.numpy() == correct:
        label = labels[i]
        prediction = predictions[i]
        image = images[i]

        ax[n].imshow(image)
        ax[n].set_title("Prediction:{}\nActual:{}".format(prediction, label))

        n += 1
        if n == (rows * cols):
          break

    if n == (rows * cols):
      break

# + colab={"base_uri": "https://localhost:8080/", "height": 214} colab_type="code" id="PSamdka2dodW" outputId="e5c521a6-3d5b-4be6-f371-a2c7a87ae124"
sample(correct=True, rows=1, cols=5)

# + [markdown] colab_type="text" id="hzHp02F_pzdh"
# Now lets take a look at where it incorrectly classifies the input. MNIST has some rather dubious handwriting, I'm sure you'll agree that some of the samples below are a little ambiguous:

# + colab={"base_uri": "https://localhost:8080/", "height": 451} colab_type="code" id="KQe5Q9LNdnb0" outputId="7f28c377-2226-468e-e388-a1cb649a26f5"
sample(correct=False, rows=2, cols=5)

StarcoderdataJupyterScriptsDedupFiltered