02_Training_DEMO.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "from numpy.random import RandomState\n",
    "from PIL import Image\n",
    "from tensorflow.keras.utils import Progbar\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.optim\n",
    "from torchvision import datasets, transforms\n",
    "import kymatio.datasets as scattering_datasets\n",
    "from contourlet_cnn import ContourletCNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Train and test\n",
    "def train(model, device, train_loader, optimizer, epoch, pbar):\n",
    "    train_loss = 0\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.cross_entropy(output, target.long())\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item()\n",
    "        pbar.update(batch_idx, values=[(\"loss\", loss.item())])\n",
    "\n",
    "            \n",
    "def test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += F.cross_entropy(output, target.long(), reduction='sum').item() # sum up batch loss\n",
    "            pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    return test_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training settings and loading training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "###### Training settings\n",
    "n_epochs = 10\n",
    "val_size = 0.2\n",
    "lr = 0.001\n",
    "seed = 2021\n",
    "img_dim = (224, 224)\n",
    "batch_size = 4\n",
    "\n",
    "# Random seed\n",
    "torch.manual_seed(seed)\n",
    "np.random.seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "use_cuda = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "\n",
    "# DataLoaders\n",
    "if use_cuda:\n",
    "    num_workers = 1\n",
    "    pin_memory = True\n",
    "else:\n",
    "    num_workers = 1\n",
    "    pin_memory = False\n",
    "\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "\n",
    "# CIFAR 10 data\n",
    "cifar_data = datasets.CIFAR10(root=scattering_datasets.get_dataset_dir('CIFAR'), train=True, transform=transforms.Compose([\n",
    "        transforms.RandomCrop(224, 224),\n",
    "        transforms.ToTensor(),\n",
    "        normalize,\n",
    "    ]), download=True)\n",
    "train_cifar_data = datasets.CIFAR10(root=scattering_datasets.get_dataset_dir('CIFAR'), train=True, transform=transforms.Compose([\n",
    "        transforms.RandomCrop(224, 224),\n",
    "        transforms.ToTensor(),\n",
    "        normalize,\n",
    "    ]), download=True)\n",
    "test_cifar_data = datasets.CIFAR10(root=scattering_datasets.get_dataset_dir('CIFAR'), train=False, transform=transforms.Compose([\n",
    "        transforms.RandomCrop(224, 224),\n",
    "        transforms.ToTensor(),\n",
    "        normalize,\n",
    "    ]), download=True)\n",
    "\n",
    "# Extract a subset of X samples per class\n",
    "prng = RandomState(seed)\n",
    "permuted = prng.permutation(np.arange(0, 5000))[0:50]\n",
    "train_random_permute = permuted[0:25]\n",
    "test_random_permute = permuted[25:50]\n",
    "train_indx = np.concatenate([np.where(np.array(cifar_data.targets) == classe)[0][train_random_permute] for classe in range(0, 10)])\n",
    "test_indx = np.concatenate([np.where(np.array(cifar_data.targets) == classe)[0][test_random_permute] for classe in range(0, 10)])\n",
    "\n",
    "train_cifar_data.data, train_cifar_data.targets = cifar_data.data[train_indx], list(np.array(cifar_data.targets)[train_indx])\n",
    "test_cifar_data.data, test_cifar_data.targets = cifar_data.data[test_indx], list(np.array(cifar_data.targets)[test_indx])\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_cifar_data,\n",
    "                                           batch_size=batch_size, shuffle=True, num_workers=num_workers,\n",
    "                                           pin_memory=pin_memory)\n",
    "test_loader = torch.utils.data.DataLoader(test_cifar_data,\n",
    "                                           batch_size=batch_size, shuffle=True, num_workers=num_workers,\n",
    "                                           pin_memory=pin_memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize Contourlet-CNN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ContourletCNN(\n",
      "  (conv_1): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (conv_2): Conv2d(88, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (conv_3): Conv2d(88, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (conv_4): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (conv_5): Conv2d(152, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (conv_6): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (conv_7): Conv2d(268, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (conv_8): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "  (conv_9): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (gap): AvgPool2d(kernel_size=7, stride=7, padding=0)\n",
      "  (fc_1): Linear(in_features=128, out_features=2048, bias=True)\n",
      "  (fc_2): Linear(in_features=2216, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = ContourletCNN(input_dim=(3, 224, 224), num_classes=10, variant=\"SSF\", spec_type=\"all\").to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters 3031386\n"
     ]
    }
   ],
   "source": [
    "model_parameters = filter(lambda p: p.requires_grad, model.parameters())\n",
    "params = sum([np.prod(p.size()) for p in model_parameters])\n",
    "print(\"Number of parameters\", params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 13.5996\n",
      "Epoch 2/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 2.5954\n",
      "Epoch 3/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 2.4112\n",
      "Epoch 4/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 2.3645\n",
      "Epoch 5/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 2.3655\n",
      "Epoch 6/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 2.3287\n",
      "Epoch 7/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 2.3338\n",
      "Epoch 8/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 2.3504\n",
      "Epoch 9/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 2.3573\n",
      "Epoch 10/10\n",
      "62/63 [============================>.] - ETA: 0s - loss: 2.3439\n",
      "920.300196266995 seconds\n"
     ]
    }
   ],
   "source": [
    "# Optimizer\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999), \n",
    "                             eps=1e-08, weight_decay=0, amsgrad=False)\n",
    "\n",
    "start_time = time.perf_counter()\n",
    "for epoch in range(n_epochs):\n",
    "    n_batches = len(train_loader)\n",
    "    pbar = Progbar(target=n_batches)\n",
    "    print(f'Epoch {epoch+1}/{n_epochs}')\n",
    "    train(model, device, train_loader, optimizer, epoch+1, pbar)\n",
    "    test(model, device, test_loader)\n",
    "    print()\n",
    "end_time = time.perf_counter()\n",
    "print(end_time - start_time, \"seconds\")"
   ]
  }
 ],
 "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}