Image Feature Extraction - AlexNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import  transforms\n",
    "import torch.nn as nn\n",
    "import torch.utils.model_zoo as model_zoo\n",
    "import glob\n",
    "import numpy as np\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# self.classifier = nn.Sequential(\n",
    "#     nn.Dropout(),\n",
    "#     nn.Linear(256 * 6 * 6, 4096),\n",
    "#     nn.ReLU(inplace=True),\n",
    "#     nn.Dropout(),\n",
    "#     nn.Linear(4096, 4096),\n",
    "#     nn.ReLU(inplace=True),\n",
    "#     nn.Linear(4096, num_classes),\n",
    "# )\n",
    "class AlexNet(nn.Module):\n",
    "\n",
    "    def __init__(self, output_layer='fc6'):\n",
    "        super(AlexNet, self).__init__()\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=3, stride=2),\n",
    "            nn.Conv2d(64, 192, kernel_size=5, padding=2),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=3, stride=2),\n",
    "            nn.Conv2d(192, 384, kernel_size=3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(384, 256, kernel_size=3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(256, 256, kernel_size=3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=3, stride=2),\n",
    "        )\n",
    "        if output_layer == 'fc6':\n",
    "            self.classifier = nn.Sequential(\n",
    "                nn.Dropout(),\n",
    "                nn.Linear(256 * 6 * 6, 4096),\n",
    "                nn.ReLU(inplace=True),\n",
    "            )\n",
    "        else:\n",
    "            assert output_layer == 'fc7'\n",
    "            self.classifier = nn.Sequential(\n",
    "                nn.Dropout(),\n",
    "                nn.Linear(256 * 6 * 6, 4096),\n",
    "                nn.ReLU(inplace=True),\n",
    "                nn.Dropout(),\n",
    "                nn.Linear(4096, 4096),\n",
    "                nn.ReLU(inplace=True),\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        x = x.view(x.size(0), 256 * 6 * 6)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "def alexnet(pretrained=False, **kwargs):\n",
    "    r\"\"\"AlexNet model architecture from the\n",
    "    `\"One weird trick...\" <https://arxiv.org/abs/1404.5997>`_ paper.\n",
    "\n",
    "    Args:\n",
    "        pretrained (bool): If True, returns a model pre-trained on ImageNet\n",
    "    \"\"\"\n",
    "    model = AlexNet(**kwargs)\n",
    "    if pretrained:\n",
    "        state = model.state_dict()\n",
    "        pretrained_weights = model_zoo.load_url('https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth')\n",
    "        assert all(k in pretrained_weights for k in state)\n",
    "        for k,v in pretrained_weights.items():\n",
    "            if k in state:\n",
    "                state[k] = v\n",
    "        model.load_state_dict(state)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "alexnet_fc6 = alexnet(pretrained=True, output_layer='fc6')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AlexNet(\n",
       "  (features): Sequential(\n",
       "    (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))\n",
       "    (1): ReLU(inplace)\n",
       "    (2): MaxPool2d(kernel_size=(3, 3), stride=(2, 2), dilation=(1, 1), ceil_mode=False)\n",
       "    (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "    (4): ReLU(inplace)\n",
       "    (5): MaxPool2d(kernel_size=(3, 3), stride=(2, 2), dilation=(1, 1), ceil_mode=False)\n",
       "    (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (7): ReLU(inplace)\n",
       "    (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (9): ReLU(inplace)\n",
       "    (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (11): ReLU(inplace)\n",
       "    (12): MaxPool2d(kernel_size=(3, 3), stride=(2, 2), dilation=(1, 1), ceil_mode=False)\n",
       "  )\n",
       "  (classifier): Sequential(\n",
       "    (0): Dropout(p=0.5)\n",
       "    (1): Linear(in_features=9216, out_features=4096, bias=True)\n",
       "    (2): ReLU(inplace)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alexnet_fc6.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "alexnet_fc7 = alexnet(pretrained=True, output_layer='fc7')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AlexNet(\n",
       "  (features): Sequential(\n",
       "    (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))\n",
       "    (1): ReLU(inplace)\n",
       "    (2): MaxPool2d(kernel_size=(3, 3), stride=(2, 2), dilation=(1, 1), ceil_mode=False)\n",
       "    (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "    (4): ReLU(inplace)\n",
       "    (5): MaxPool2d(kernel_size=(3, 3), stride=(2, 2), dilation=(1, 1), ceil_mode=False)\n",
       "    (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (7): ReLU(inplace)\n",
       "    (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (9): ReLU(inplace)\n",
       "    (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (11): ReLU(inplace)\n",
       "    (12): MaxPool2d(kernel_size=(3, 3), stride=(2, 2), dilation=(1, 1), ceil_mode=False)\n",
       "  )\n",
       "  (classifier): Sequential(\n",
       "    (0): Dropout(p=0.5)\n",
       "    (1): Linear(in_features=9216, out_features=4096, bias=True)\n",
       "    (2): ReLU(inplace)\n",
       "    (3): Dropout(p=0.5)\n",
       "    (4): Linear(in_features=4096, out_features=4096, bias=True)\n",
       "    (5): ReLU(inplace)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alexnet_fc7.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RemoveAlphaChannel(object):\n",
    "    def __call__(self, image_tensor):        \n",
    "        if image_tensor.shape[0] == 4:\n",
    "            image_tensor = image_tensor[:3]\n",
    "        assert image_tensor.shape[0] == 3\n",
    "        return image_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImageDataset(Dataset):\n",
    "    def __init__(self):\n",
    "        self.image_files = glob.glob(\"/mnt/workspace/Ugallery/images/*.jpg\")        \n",
    "        self.artwork_ids = [int(f[f.rfind('/')+1:-4]) for f in self.image_files]\n",
    "        self.transform = transforms.Compose([\n",
    "#             transforms.Scale(256),\n",
    "#             transforms.CenterCrop(224),\n",
    "            transforms.Scale([224,224]),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(\n",
    "                mean=[0.485, 0.456, 0.406],\n",
    "                std=[0.229, 0.224, 0.225]\n",
    "            ),\n",
    "            RemoveAlphaChannel()\n",
    "        ])\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_files)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {'id': self.artwork_ids[idx],\n",
    "                'image': self.transform(Image.open(self.image_files[idx]))}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_dataset = ImageDataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = DataLoader(image_dataset, batch_size=1024, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = len(image_dataset)\n",
    "output_ids = np.empty((N,), dtype=int)\n",
    "output_fc6 = np.empty((N,4096), dtype=float)\n",
    "output_fc7 = np.empty((N,4096), dtype=float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i_batch = 0, offset = 1024, batch_size = 1024, elapsed_time = 39.50\n",
      "i_batch = 1, offset = 2048, batch_size = 1024, elapsed_time = 68.87\n",
      "i_batch = 2, offset = 3072, batch_size = 1024, elapsed_time = 98.30\n",
      "i_batch = 3, offset = 4096, batch_size = 1024, elapsed_time = 128.13\n",
      "i_batch = 4, offset = 5120, batch_size = 1024, elapsed_time = 157.83\n",
      "i_batch = 5, offset = 6144, batch_size = 1024, elapsed_time = 186.23\n",
      "i_batch = 6, offset = 7168, batch_size = 1024, elapsed_time = 214.57\n",
      "i_batch = 7, offset = 8192, batch_size = 1024, elapsed_time = 242.88\n",
      "i_batch = 8, offset = 9216, batch_size = 1024, elapsed_time = 271.32\n",
      "i_batch = 9, offset = 10240, batch_size = 1024, elapsed_time = 299.65\n",
      "i_batch = 10, offset = 11264, batch_size = 1024, elapsed_time = 327.93\n",
      "i_batch = 11, offset = 12288, batch_size = 1024, elapsed_time = 356.21\n",
      "i_batch = 12, offset = 13297, batch_size = 1009, elapsed_time = 384.07\n"
     ]
    }
   ],
   "source": [
    "offset = 0\n",
    "start_time = time()\n",
    "for i_batch, batch in enumerate(dataloader):    \n",
    "    assert offset < N\n",
    "    \n",
    "    batch_ids = batch['id'].numpy()\n",
    "    batch_var = Variable(batch['image'])\n",
    "    batch_size = len(batch_ids)\n",
    "    \n",
    "    batch_fc6 = alexnet_fc6(batch_var).data.numpy()\n",
    "    batch_fc7 = alexnet_fc7(batch_var).data.numpy()\n",
    "    \n",
    "    output_ids[offset:offset+batch_size] = batch_ids\n",
    "    output_fc6[offset:offset+batch_size] = batch_fc6\n",
    "    output_fc7[offset:offset+batch_size] = batch_fc7\n",
    "    \n",
    "    elapsed_time = time() - start_time    \n",
    "    offset += batch_size\n",
    "    \n",
    "    print('i_batch = %d, offset = %d, batch_size = %d, elapsed_time = %.2f' % (i_batch, offset, batch_size, elapsed_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.8294545454770808, 0.2906904877248657)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_fc6.mean(), output_fc7.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "dir_path = \"/mnt/workspace/Ugallery/AlexNet/\"\n",
    "os.makedirs(dir_path, exist_ok=True)\n",
    "output_fc6.dump(dir_path + \"fc6.npy\")\n",
    "output_fc7.dump(dir_path + \"fc7.npy\")\n",
    "output_ids.dump(dir_path + \"ids.npy\")"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}