train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7507140",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import datetime\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.backends.cudnn as cudnn\n",
    "\n",
    "from timm.models import create_model\n",
    "\n",
    "from engine import train_one_epoch, train_one_epoch_distillation, evaluate\n",
    "from utils import get_training_dataloader, get_test_dataloader\n",
    "import models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "234c6ce8",
   "metadata": {},
   "outputs": [],
   "source": [
    "MEAN = (0.5070751592371323, 0.48654887331495095, 0.4409178433670343)\n",
    "STD = (0.2673342858792401, 0.2564384629170883, 0.27615047132568404)\n",
    "CHECKPOINT_PATH = './checkpoint'\n",
    "MODEL_NAME = 'vit_tiny_patch16_224'\n",
    "num_classes = 10\n",
    "EPOCHS = 5\n",
    "LR = 0.0001\n",
    "WD = 0.0\n",
    "shots = 1000\n",
    "\n",
    "print(f\"Creating model: {MODEL_NAME}\")\n",
    "model = create_model(\n",
    "        MODEL_NAME,\n",
    "        pretrained=True,\n",
    "        num_classes=10,\n",
    "        img_size=224)\n",
    "device = 'cuda:0' # device = 'cpu'\n",
    "model = model.to(device)\n",
    "\n",
    "cifar10_training_loader = get_training_dataloader(\n",
    "    MEAN,\n",
    "    STD,\n",
    "    num_workers=2,\n",
    "    batch_size=16,\n",
    "    shuffle=True,\n",
    "    shots=shots\n",
    ")\n",
    "\n",
    "assert (shots*num_classes == len(cifar10_training_loader.dataset))\n",
    "\n",
    "cifar10_test_loader = get_test_dataloader(\n",
    "    MEAN,\n",
    "    STD,\n",
    "    num_workers=4,\n",
    "    batch_size=256,\n",
    "    shuffle=False\n",
    ")\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=WD)\n",
    "\n",
    "\n",
    "n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print('number of params:', n_parameters)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f51e2794",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Start training for {EPOCHS} epochs\")\n",
    "\n",
    "for epoch in range(1, EPOCHS+1):\n",
    "    train_stats = train_one_epoch(\n",
    "        model, criterion, cifar10_training_loader,\n",
    "        optimizer, device, epoch)\n",
    "    if epoch % 10 == 9: \n",
    "        test_stats = evaluate(cifar10_test_loader, model, criterion, device)\n",
    "        print(f\"Accuracy of the network on the {len(cifar10_test_loader)} test images: {test_stats['acc1']:.1f}%\")        \n",
    "        \n",
    "test_stats = evaluate(cifar10_test_loader, model, criterion, device)\n",
    "print(f\"Accuracy of the network on the {len(cifar10_test_loader)} test images: {test_stats['acc1']:.1f}%\")            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e323d040",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate througput \n",
    "start_time = time.time()\n",
    "test_stats = evaluate(cifar10_test_loader, model, criterion, device)\n",
    "end_time = time.time()\n",
    "num_samples = len(cifar10_test_loader.dataset)\n",
    "throughput = num_samples / (end_time - start_time)\n",
    "print(\"Throughput: {}\".format(throughput))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caaf4f8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Train the teacher\n",
    "\n",
    "MODEL_NAME = 'vit_base_patch16_224'\n",
    "num_classes = 10\n",
    "EPOCHS = 5\n",
    "LR = 0.0001\n",
    "WD = 0.0\n",
    "\n",
    "print(f\"Creating model: {MODEL_NAME}\")\n",
    "teacher = create_model(\n",
    "        MODEL_NAME,\n",
    "        pretrained=True,\n",
    "        num_classes=10,\n",
    "        img_size=224)\n",
    "device = 'cuda:0' # device = 'cpu'\n",
    "teacher = teacher.to(device)\n",
    "\n",
    "\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(teacher.parameters(), lr=LR, weight_decay=WD)\n",
    "\n",
    "\n",
    "n_parameters = sum(p.numel() for p in teacher.parameters() if p.requires_grad)\n",
    "print('number of params:', n_parameters)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c392f94",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Start training for {EPOCHS} epochs\")\n",
    "\n",
    "for epoch in range(1, EPOCHS+1):\n",
    "    train_stats = train_one_epoch(\n",
    "        teacher, criterion, cifar10_training_loader,\n",
    "        optimizer, device, epoch)\n",
    "    if epoch % 10 == 9: \n",
    "        test_stats = evaluate(cifar10_test_loader, teacher, criterion, device)\n",
    "        print(f\"Accuracy of the network on the {len(cifar10_test_loader)} test images: {test_stats['acc1']:.1f}%\")        \n",
    "        \n",
    "test_stats = evaluate(cifar10_test_loader, teacher, criterion, device)\n",
    "print(f\"Accuracy of the network on the {len(cifar10_test_loader)} test images: {test_stats['acc1']:.1f}%\")       \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab509fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save finetuned teacher model\n",
    "torch.save(teacher.state_dict(), './teacher.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69587aa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "teacher = create_model(\n",
    "        'vit_base_patch16_224',\n",
    "        pretrained=True,\n",
    "        num_classes=10,\n",
    "        img_size=224)\n",
    "device = 'cuda:0' # device = 'cpu'\n",
    "teacher = teacher.to(device)\n",
    "teacher.load_state_dict(torch.load('./teacher.pth'))\n",
    "\n",
    "test_stats = evaluate(cifar10_test_loader, teacher, criterion, device)\n",
    "print(f\"Accuracy of the network on the {len(cifar10_test_loader)} test images: {test_stats['acc1']:.1f}%\")            \n",
    "\n",
    "# Train the student\n",
    "for p in teacher.parameters(): \n",
    "    p.requires_grad = False\n",
    "\n",
    "MODEL_NAME = 'vit_tiny_patch16_224'\n",
    "\n",
    "model = create_model(\n",
    "        MODEL_NAME,\n",
    "        pretrained=True,\n",
    "        num_classes=10,\n",
    "        img_size=224)\n",
    "device = 'cuda:0' # device = 'cpu'\n",
    "model = model.to(device)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=WD)\n",
    "\n",
    "\n",
    "n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print('number of params:', n_parameters)\n",
    "\n",
    "\n",
    "print(f\"Start training for {EPOCHS} epochs\")\n",
    "\n",
    "for epoch in range(1, EPOCHS+1):\n",
    "    train_stats = train_one_epoch_distillation(\n",
    "        teacher, model, criterion, cifar10_training_loader,\n",
    "        optimizer, device, epoch, alpha=2.0, temp=1.0)\n",
    "    if epoch % 2 == 1: \n",
    "        test_stats = evaluate(cifar10_test_loader, model, criterion, device)\n",
    "        print(f\"Accuracy of the network on the {len(cifar10_test_loader)} test images: {test_stats['acc1']:.1f}%\")        \n",
    "        \n",
    "test_stats = evaluate(cifar10_test_loader, model, criterion, device)\n",
    "print(f\"Accuracy of the network on the {len(cifar10_test_loader)} test images: {test_stats['acc1']:.1f}%\")            "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}