pytorch/computer_vision/lightweight_networks_and_mobileNet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have seen that complex network require significant resources, such as GPU, for training, and also for fast inference. However, it turns out that a model with significantly smaller number of parameters in most cases can still be trained to perform reasonably well. In other worlds, increase in the model complexity typically results in small(non-proportional) increase in the model performance.\n",
    "\n",
    "According the previously notebooks, we can see that the accuracy of simple dense model was not significantly worse than that of a poweful CNN. **Increasing the number of CNN layer and/or number of neurons in the classifier allowed us to gain a few percents of accuracy at most**.\n",
    "\n",
    "This leads us to the idea that we can experiment with `Lightweight network architectures` in order to train faster models. This is especially important if we want to be able to execute our models on mobile devices.\n",
    "\n",
    "This module will rely on the Cats and Dogs dataset. First we will make sure that the dataset is available.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "from torchinfo import summary\n",
    "import os, glob, zipfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check the paltform, Apple Silicon or Linux\n",
    "import os, platform\n",
    "\n",
    "torch_device=\"cpu\"\n",
    "\n",
    "if 'kaggle' in os.environ.get('KAGGLE_URL_BASE','localhost'):\n",
    "    torch_device = 'cuda'\n",
    "else:\n",
    "    torch_device = 'mps' if platform.system() == 'Darwin' else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('data/kagglecatsanddogs_5340.zip'):\n",
    "    !wget -P data https://download.microsoft.com/download/3/E/1/3E1C3F21-ECDB-4869-8368-6DEBA77B919F/kagglecatsanddogs_5340.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    "def check_image(fn):\n",
    "    try:\n",
    "        im = Image.open(fn)\n",
    "        im.verify()\n",
    "        return True\n",
    "    except:\n",
    "        return False\n",
    "\n",
    "def check_image_dir(path):\n",
    "    for fn in glob.glob(path):\n",
    "        if not check_image(fn):\n",
    "            print(\"Corrupt image: {}\".format(fn))\n",
    "            os.remove(fn)\n",
    "\n",
    "def common_transform():\n",
    "    # torchvision.transforms.Normalize is used to normalize a tensor image with mean and standard deviation.\n",
    "    std_normalize = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                                     std=[0.229, 0.224, 0.225])\n",
    "    # torchvision.transforms.Compose is used to compose several transforms together in order to do data augmentation.\n",
    "    trans = torchvision.transforms.Compose([\n",
    "        torchvision.transforms.Resize(256), # resize the image to 256x256\n",
    "        torchvision.transforms.CenterCrop(224), # crop the image to 224x224 about the center\n",
    "        torchvision.transforms.ToTensor(), # convert the image to a tensor with pixel values in the range [0, 1]\n",
    "        std_normalize])\n",
    "    return trans\n",
    "\n",
    "def load_cats_dogs_dataset():\n",
    "    if not os.path.exists('data/PetImages'):\n",
    "        with zipfile.ZipFile('data/kagglecatsanddogs_5340.zip', 'r') as zip_ref:\n",
    "            zip_ref.extractall('data')\n",
    "    \n",
    "    check_image_dir('data/PetImages/Cat/*.jpg')\n",
    "    check_image_dir('data/PetImages/Dog/*.jpg')\n",
    "\n",
    "    dataset = torchvision.datasets.ImageFolder('data/PetImages', transform=common_transform())\n",
    "    trainset, testset = torch.utils.data.random_split(dataset, [20000, len(dataset) - 20000])\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True, num_workers=2) # num_workers: how many subprocesses to use for data loading\n",
    "    testloader = torch.utils.data.DataLoader(testset, batch_size=32, shuffle=False, num_workers=2)\n",
    "    return dataset, trainloader, testloader \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset, trainloader, testloader = load_cats_dogs_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MobileNet\n",
    "\n",
    "In the previous notebook, we habve seen [**ResNet** architecture](https://www.kaggle.com/code/aisuko/pre-trained-models-and-transfer-learning) for image classification. More lightweight analog of ResNet is **MobileNet**, which uses so-called *Inverted Residual Blocks*. Let's load pre-trained mobilenet and see how it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(torch.__version__)\n",
    "# https://pytorch.org/hub/pytorch_vision_mobilenet_v2/\n",
    "model = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=True)\n",
    "model.eval()\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Apply the model to the dataset and visualize the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_image = dataset[0][0].unsqueeze(0) # unsqueeze(0): add a dimension of size 1 at the 0th position\n",
    "res = model(sample_image) # apply the model to the sample image\n",
    "print(res[0].argmax()) # get the index of the highest probability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using MobileNet for transfer learning\n",
    "\n",
    "Now let's perform the same transfer learning process as in previous notebook, but using MobileNet as a base model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Freeze all parameters of the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for x in model.parameters():\n",
    "    x.requires_grad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Replace the final classifier\n",
    "\n",
    "We also transfer the model to our default training device (GPU or CPU)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check the paltform, Apple Silicon or Linux\n",
    "import os, platform\n",
    "\n",
    "torch_device=\"cpu\"\n",
    "\n",
    "if 'kaggle' in os.environ.get('KAGGLE_URL_BASE','localhost'):\n",
    "    torch_device = 'cuda'\n",
    "else:\n",
    "    torch_device = 'mps' if platform.system() == 'Darwin' else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.classifier = nn.Linear(1280,2)  # change the last layer to a linear layer with 2 outputs\n",
    "model = model.to(torch_device)\n",
    "summary(model, input_size=(1, 3, 224, 224))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Doing the actual training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(net, dataloader, loss_fn=nn.NLLLoss()):\n",
    "    net.eval() # put the network into evaluation mode to deactivate the dropout layers\n",
    "    count,acc,loss =0,0,0\n",
    "    with torch.no_grad(): # deactivate autograd to save memory and speed up computations\n",
    "        for features, labels in dataloader:\n",
    "            features,labels = features.to(torch_device), labels.to(torch_device)\n",
    "            out=net(features) # forward pass of the mini-batch through the network to obtain the outputs\n",
    "            loss += loss_fn(out,labels) # compute the loss\n",
    "            preds=torch.max(out,dim=1)[1] # compute the predictions to obtain the accuracy\n",
    "            acc+=(preds==labels).sum() # accumulate the correct predictions\n",
    "            count+=len(labels) # accumulate the total number of examples\n",
    "    return loss.item()/count, acc.item()/count # return the loss and accuracy\n",
    "\n",
    "def train_long(net, train_loader, test_loader, epochs=5, lr=0.01, optimizer=None, loss_fn=nn.NLLLoss(), print_freq=10):\n",
    "    optimizer = optimizer or torch.optim.Adam(net.parameters(), lr=lr) # use Adam optimizer if not provided\n",
    "    for epoch in range(epochs):\n",
    "        net.train() # put the network into training mode make sure the parameters are trainable\n",
    "        total_loss,acc,count =0,0,0\n",
    "        for i, (features, labels) in enumerate(train_loader):\n",
    "            lbls = labels.to(torch_device)\n",
    "            optimizer.zero_grad() # reset the gradients to zero before each batch to avoid accumulation\n",
    "            out=net(features.to(torch_device)) # forward pass of the mini-batch through the network to obtain the outputs\n",
    "            loss = loss_fn(out, lbls) # compute the loss\n",
    "            loss.backward() # compute the gradients of the loss with respect to all the parameters of the network\n",
    "            optimizer.step() # update the parameters of the network using the gradients to minimize the loss\n",
    "            total_loss+=loss # accumulate the loss for inspection\n",
    "            _,preds=torch.max(out,dim=1) # compute the predictions to obtain the accuracy\n",
    "            acc+=(preds==lbls).sum() # accumulate the correct predictions\n",
    "            count+=len(lbls) # accumulate the total number of examples\n",
    "            if i%print_freq==0:\n",
    "                print(f'Epoch {epoch}, iter {i}, loss={total_loss.item()/count:.3f}, acc={acc.item()/count:.3f}')\n",
    "        vl, va = validate(net, test_loader, loss_fn=loss_fn)\n",
    "        print(f'Epoch {epoch}, val_loss={vl:.3f}, val_acc={va:.3f}')\n",
    "\n",
    "train_long(model, trainloader, testloader, loss_fn=torch.nn.CrossEntropyLoss(),epochs=1, print_freq=90)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Notice that MobileNet results in almost the same accuracy as VGG-16, and just slightly lower than full-scale ResNet.\n",
    "\n",
    "The main advantage of small models, such as MobileNet or ResNet-18 is that they can be used on mobile devices, "
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}