{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Guide to Pretraining\n",
    "In this guide, you'll learn how to pretrain a model without using labels.\n",
    "\n",
    "First, import Masterful alongside the other necessary libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "import sys\n",
    "sys.path.append('../../../')\n",
    "\n",
    "import masterful\n",
    "import masterful.backend.pretrain as pretrain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this guide we'll be using the CIFAR10 dataset, which consists of 60,000 color images of 10 separate, non-overlapping classes of objects. There are 6,000 images of each class of object; 5,000 in a \"training\" set and 1,000 held out for testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
    "validation_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a Data Specification, which captures CIFAR10's relevant metadata:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_spec = masterful.DataSpec.from_dataset(\n",
    "        task=masterful.spec.Task.CLASSIFICATION,\n",
    "        dataset=train_dataset,\n",
    "        image_range=masterful.spec.ImageRange.ZERO_255,\n",
    "        num_classes=10,\n",
    "        sparse=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now create a Pretraining Policy, which establishes pretraining-related hyperparameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrain_policy = masterful.PretrainPolicy(batch_size=512,\n",
    "                                           epochs=3,\n",
    "                                           warmup_steps=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, create a model to be pretrained – here, the same simple CNN used in [this](https://www.tensorflow.org/tutorials/images/cnn) TensorFlow tutorial for labeled training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = tf.keras.models.Sequential()\n",
    "model.add(tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)))\n",
    "model.add(tf.keras.layers.MaxPooling2D((2, 2)))\n",
    "model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))\n",
    "model.add(tf.keras.layers.MaxPooling2D((2, 2)))\n",
    "model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you're ready to pretrain! Pretraining with Masterful will return a training report, including fields like\n",
    "- `loss`: the unsupervised loss at the end of pretraining.\n",
    "- `accuracy`: the accuracy achieved by performing K-Nearest Neighbors classification on the pretrained backbone's output features\n",
    "\n",
    "(Note that the labels in the dataset passed to the `labeled_data` parameter below are only used for KNN classification, _not_ for pretraining the backbone.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/3\n",
      "98/98 [==============================] - 155s 2s/step - loss: 1159.5908\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Feature extracting: 100%|██████████| 98/98 [00:00<00:00, 124.99it/s]\n",
      "Test Epoch: Acc@1:32.01%: 100%|██████████| 20/20 [00:00<00:00, 37.86it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kNN Test Accuracy at epoch 0: 32.01000213623047 Max Accuracy so far: 32.01000213623047\n",
      "Epoch 2/3\n",
      "98/98 [==============================] - 153s 2s/step - loss: 982.5400\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Feature extracting: 100%|██████████| 98/98 [00:00<00:00, 132.23it/s]\n",
      "Test Epoch: Acc@1:34.45%: 100%|██████████| 20/20 [00:00<00:00, 40.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kNN Test Accuracy at epoch 1: 34.45000076293945 Max Accuracy so far: 34.45000076293945\n",
      "Epoch 3/3\n",
      "98/98 [==============================] - 153s 2s/step - loss: 923.1833\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Feature extracting: 100%|██████████| 98/98 [00:00<00:00, 135.27it/s]\n",
      "Test Epoch: Acc@1:35.30%: 100%|██████████| 20/20 [00:00<00:00, 37.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kNN Test Accuracy at epoch 2: 35.29999923706055 Max Accuracy so far: 35.29999923706055\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "training_report = pretrain.pretrain(\n",
    "        pretrain_policy=pretrain_policy,\n",
    "        model=model,\n",
    "        validation_data=validation_dataset,\n",
    "        labeled_data=train_dataset,\n",
    "        unlabeled_data=None,\n",
    "        synthetic_data=None,\n",
    "        data_spec=data_spec,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pretraining policy above only runs for a few epochs, to save time; you should expect better results with more epochs (and quicker results for a smaller model). However, note that the pretrained model's output features—without the use of labels—already outperform the guessing of randomly initialized weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final pretraining loss: 923.1832885742188\n",
      "Final pretraining KNN accuracy: 35.29999923706055\n"
     ]
    }
   ],
   "source": [
    "loss = training_report.validation_metrics['loss']\n",
    "acc = training_report.validation_metrics['accuracy']\n",
    "print(f'Final pretraining loss: {loss}')\n",
    "print(f'Final pretraining KNN accuracy: {acc}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you're ready to use Masterful's unsupervised pretraining API!"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}