semisupervised_classification_keras.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "semisupervised-classification-keras.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gSRJBo9-lwOX"
      },
      "source": [
        "# Semi-supervised Classification\n",
        "\n",
        "This jupyter notebook contains a training script for the https://github.com/beresandras/semisupervised-classification-keras repository, and is intended to be used in a Google Colab environment."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1-Q2l8ZretUZ"
      },
      "source": [
        "# comment out after first run\n",
        "!git clone https://github.com/beresandras/semisupervised-classification-keras"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mKsXIajqePwJ"
      },
      "source": [
        "import sys\n",
        "import pickle\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow.keras.layers.experimental import preprocessing\n",
        "\n",
        "sys.path.insert(0,'/content/semisupervised-classification-keras')\n",
        "from dataset import prepare_dataset\n",
        "from augmentations import RandomResizedCrop, RandomColorJitter\n",
        "from algorithms import CrossEntropy, InfoNCE, SuNCEt, PAWS\n",
        "from models import KNNClassifier"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rAXcug3LeU4P"
      },
      "source": [
        "# hyperparameters\n",
        "num_epochs = 25\n",
        "width = 128\n",
        "k_values = (20, 200) # for knn classification\n",
        "batch_sizes = {  # unlabeled, labeled\n",
        "    CrossEntropy: (250, 250),\n",
        "    InfoNCE: (500, 25),\n",
        "    SuNCEt: (250, 250),\n",
        "    PAWS: (250, 250),\n",
        "}\n",
        "hyperparams = {\n",
        "    CrossEntropy: {},\n",
        "    InfoNCE: {\"temperature\": 0.1},\n",
        "    SuNCEt: {\"temperature\": 0.1, \"supervised_loss_weight\": 1.0},\n",
        "    PAWS: {\"temperature\": 0.1, \"sharpening\": 0.25},\n",
        "}\n",
        "\n",
        "# select an algorithm\n",
        "Algorithm = SuNCEt  # CrossEntropy, InfoNCE, SuNCEt, PAWS"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xKX-vD2jeVLe"
      },
      "source": [
        "# load STL10 dataset\n",
        "train_dataset, eval_train_dataset, eval_test_dataset = prepare_dataset(\n",
        "    *batch_sizes[Algorithm]\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g3NnbOlWeVXN"
      },
      "source": [
        "model = Algorithm(\n",
        "    augmenter=keras.Sequential(\n",
        "        [\n",
        "            layers.Input(shape=(96, 96, 3)),\n",
        "            preprocessing.Rescaling(1 / 255),\n",
        "            preprocessing.RandomFlip(\"horizontal\"),\n",
        "            RandomResizedCrop(scale=(0.2, 1.0), ratio=(3 / 4, 4 / 3)),\n",
        "            RandomColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.2),\n",
        "        ],\n",
        "        name=\"augmenter\",\n",
        "    ),\n",
        "    encoder=keras.Sequential(\n",
        "        [\n",
        "            layers.Input(shape=(96, 96, 3)),\n",
        "            layers.Conv2D(width, kernel_size=3, strides=2, activation=\"relu\"),\n",
        "            layers.Conv2D(width, kernel_size=3, strides=2, activation=\"relu\"),\n",
        "            layers.Conv2D(width, kernel_size=3, strides=2, activation=\"relu\"),\n",
        "            layers.Conv2D(width, kernel_size=3, strides=2, activation=\"relu\"),\n",
        "            layers.Flatten(),\n",
        "            layers.Dense(width),\n",
        "        ],\n",
        "        name=\"encoder\",\n",
        "    ),\n",
        "    **hyperparams[Algorithm]\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kCBVd6DBeViQ"
      },
      "source": [
        "evaluators = [\n",
        "    KNNClassifier(\n",
        "        train_dataset=eval_train_dataset,\n",
        "        test_dataset=eval_test_dataset,\n",
        "        model=model,\n",
        "        k=k,\n",
        "    )\n",
        "    for k in k_values\n",
        "]\n",
        "model.compile(\n",
        "    optimizer=keras.optimizers.Adam(),\n",
        ")\n",
        "\n",
        "# run training\n",
        "history = model.fit(\n",
        "    train_dataset,\n",
        "    epochs=num_epochs,\n",
        "    validation_data=eval_test_dataset,\n",
        "    callbacks=[\n",
        "        keras.callbacks.LambdaCallback(on_epoch_end=evaluator.evaluate)\n",
        "        for evaluator in evaluators\n",
        "    ],\n",
        ")\n",
        "\n",
        "# save history\n",
        "with open(\"{}.pkl\".format(Algorithm.__name__), \"wb\") as output_file:\n",
        "    pickle.dump(history.history, output_file)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}