Train_puzzle_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import math\n",
    "import types\n",
    "import random\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.data as Data\n",
    "from torch.autograd import Variable\n",
    "from torchvision import datasets, transforms, models\n",
    "from sklearn.metrics import classification_report\n",
    "from torch import optim\n",
    "import time\n",
    "from utils import Conv2d\n",
    "from torch.utils.data import Dataset, DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters setting\n",
    "num_classes = 7\n",
    "cut_size = 225\n",
    "init_lr = 0.001\n",
    "batch_size = 64\n",
    "device= torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data\n",
    "data_dir = ''\n",
    "\n",
    "train_data = torch.load(data_dir + 'pickles/AffectNet_train_full_p1.pt')\n",
    "train_img, train_y = train_data['images'], train_data['labels']\n",
    "\n",
    "test_data = torch.load(data_dir + 'pickles/AffectNet_test.pt')\n",
    "test_img, test_y = test_data['images'], test_data['emotion']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nine_crops(image, size=[96,96], n_crop=3):\n",
    "    xx = yy = n_crop\n",
    "    x = size[0] // xx\n",
    "    y = size[1] // yy\n",
    "    crops = ()\n",
    "    for j in range(yy):\n",
    "        for i in range(xx):\n",
    "            left = i*x\n",
    "            up = y*j\n",
    "            right = left + x\n",
    "            low = up + y\n",
    "            region = image.crop((left,up,right,low))\n",
    "            region = transforms.Resize([cut_size//n_crop,cut_size//n_crop])(region)\n",
    "            crops += (region,)       \n",
    "    return crops\n",
    "\n",
    "class puzzle_crop(object):\n",
    "    def __init__(self, size, n_crop):\n",
    "        self.size = size\n",
    "        self.n_crop = n_crop\n",
    "\n",
    "    def __call__(self, img):\n",
    "        return nine_crops(img, self.size, self.n_crop)\n",
    "\n",
    "\n",
    "class Lambda(object):\n",
    "    def __init__(self, lambd):\n",
    "        assert isinstance(lambd, types.LambdaType)\n",
    "        self.lambd = lambd\n",
    "    def __call__(self, img):\n",
    "        return self.lambd(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "#======= data loaders =======#    \n",
    "n_crop = 3\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.ToPILImage(),\n",
    "    transforms.Grayscale(num_output_channels=3),\n",
    "    transforms.Resize([int(cut_size*1.1),int(cut_size*1.1)]),\n",
    "    transforms.RandomCrop([cut_size,cut_size]),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    puzzle_crop([cut_size,cut_size],3),\n",
    "    Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops]))\n",
    "])\n",
    "\n",
    "test_transform = transforms.Compose([\n",
    "    transforms.ToPILImage(),\n",
    "    transforms.Grayscale(num_output_channels=3),\n",
    "    puzzle_crop([cut_size,cut_size],3),\n",
    "    Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops]))\n",
    "])\n",
    "\n",
    "resize_transform = transforms.Compose([\n",
    "    transforms.ToPILImage(),\n",
    "    transforms.Grayscale(num_output_channels=3),\n",
    "    transforms.Resize([cut_size,cut_size]),\n",
    "    transforms.ToTensor()\n",
    "])\n",
    "\n",
    "def default_loader(images):\n",
    "    img_tensor = train_transform(images)\n",
    "    return img_tensor\n",
    "\n",
    "def default_test_loader(images):\n",
    "    img_tensor = test_transform(images)\n",
    "    return img_tensor\n",
    "\n",
    "class Dataset(Dataset):\n",
    "    def __init__(self,imgs,loader=default_loader):\n",
    "        self.images = imgs \n",
    "        self.loader = loader\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "\n",
    "        img = self.images[index]\n",
    "        n_imgs = tmp = self.loader(img)\n",
    "        img = resize_transform(img)\n",
    "        label = torch.randperm(n_crop**2)\n",
    "        for i in range(label.size(0)):\n",
    "            n_imgs[i] = tmp[label[i]]\n",
    "        return img, n_imgs,label\n",
    "\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PairPuzzle_AE(nn.Module):\n",
    "\n",
    "    def __init__(self,model=None,lr=0.0001,num_classes=1000):\n",
    "        super(PairPuzzle_AE, self).__init__()\n",
    "        \n",
    "        self.feature = nn.Sequential(\n",
    "            # conv_0\n",
    "            Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(32),\n",
    "            # conv_1\n",
    "            Conv2d(in_channels=32, out_channels=64, kernel_size=2, stride=2),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(64),\n",
    "            # conv_2\n",
    "            Conv2d(in_channels=64, out_channels=128, kernel_size=2, stride=2),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(128),\n",
    "            # conv_3\n",
    "            Conv2d(in_channels=128, out_channels=256, kernel_size=2, stride=2),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(256),\n",
    "            # conv_4\n",
    "            Conv2d(in_channels=256, out_channels=512, kernel_size=2, stride=2),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(512)\n",
    "        )\n",
    "        \n",
    "        self.fc6 = nn.Sequential()\n",
    "        self.fc6.add_module('fc6_s1',nn.Linear(12800, 4608))\n",
    "        self.fc6.add_module('relu6_s1',nn.ReLU(inplace=True))\n",
    "        self.fc6.add_module('drop6_s1',nn.Dropout(p=0.5))\n",
    "        self.classifier = nn.Sequential(\n",
    "#             nn.Linear(9472,4608),\n",
    "#             nn.ReLU(),\n",
    "            nn.Linear(4608,4096),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=0.5),\n",
    "            nn.Linear(4096,num_classes)\n",
    "        )\n",
    "        \n",
    "        self.loss_fn = nn.CrossEntropyLoss() #FocalLoss(4.0)\n",
    "#         self.optimizer=torch.optim.Adam([p for p in self.parameters() if p.requires_grad] , lr=lr)\n",
    "        self.optimizer=torch.optim.SGD(self.parameters(),lr=lr,momentum=0.9,weight_decay = 5e-4)\n",
    "\n",
    "\n",
    "    def forward(self,x,n_x):\n",
    "        B,T,C,H,W = n_x.size()\n",
    "        x = self.feature(x).view(B,1,-1)\n",
    "        n_x = n_x.transpose(0,1)\n",
    "        x_list = []\n",
    "        for i in range(9):\n",
    "            z = self.feature(n_x[i])\n",
    "            z = self.fc6(z.view(B,-1))\n",
    "            z = z.view([B,1,-1])\n",
    "            x = torch.cat((z,x),-1)\n",
    "            x_list.append(z)\n",
    "        x = torch.cat(x_list,1).view(B*T,-1)\n",
    "        p = self.classifier(x).view(B,T,-1)\n",
    "        \n",
    "        return x,p\n",
    "        \n",
    "    def train_model(self, data_loader, num_epochs, lr, min_loss):\n",
    "        self.train()\n",
    "        start=time.time()\n",
    "\n",
    "        for epoch in range(num_epochs):\n",
    "            running_loss=0\n",
    "            running_acc = 0.\n",
    "            total = 0\n",
    "            num_batches = 0\n",
    "            for count, (images, n_images, label) in enumerate(data_loader):\n",
    "                self.optimizer.zero_grad()\n",
    "                \n",
    "                images=images.to(device)\n",
    "                n_images=n_images.to(device)\n",
    "                label=label.to(device)\n",
    "\n",
    "                features, scores = self.forward(images, n_images)\n",
    "\n",
    "                loss =  self.loss_fn(scores, label) \n",
    "                loss.backward()\n",
    "\n",
    "                self.optimizer.step()\n",
    "\n",
    "                running_loss += loss.detach().item()\n",
    "\n",
    "                acc = (scores.argmax(-1) == label).float().sum()\n",
    "                running_acc += acc.item()/9\n",
    "\n",
    "                total+=images.size(0)\n",
    "                num_batches += 1\n",
    "            total_loss = running_loss/num_batches\n",
    "            if total_loss<min_loss:\n",
    "                min_loss = total_loss\n",
    "                print(min_loss)\n",
    "                torch.save(model.state_dict(),data_dir+'BaselinesGray/CNN_Puzzle_GrayAffectNet_inputsize{}_minLoss.pth'.format(cut_size))\n",
    "            total_acc = running_acc/total\n",
    "            elapsed = time.time()-start\n",
    "            print('\\t', 'epoch=',epoch, '\\t time = {:.0f}m {:.0f}s'.format(elapsed // 60, elapsed % 60),'\\t lr=', lr,'\\t loss = ', total_loss , '\\t Acc = {:.6f}'.format(total_acc*100),'%')\n",
    "        return min_loss\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PairPuzzle_AE(\n",
      "  (feature): Sequential(\n",
      "    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1))\n",
      "    (1): ReLU()\n",
      "    (2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (3): Conv2d(32, 64, kernel_size=(2, 2), stride=(2, 2))\n",
      "    (4): ReLU()\n",
      "    (5): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (6): Conv2d(64, 128, kernel_size=(2, 2), stride=(2, 2))\n",
      "    (7): ReLU()\n",
      "    (8): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (9): Conv2d(128, 256, kernel_size=(2, 2), stride=(2, 2))\n",
      "    (10): ReLU()\n",
      "    (11): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (12): Conv2d(256, 512, kernel_size=(2, 2), stride=(2, 2))\n",
      "    (13): ReLU()\n",
      "    (14): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (fc6): Sequential(\n",
      "    (fc6_s1): Linear(in_features=12800, out_features=4608, bias=True)\n",
      "    (relu6_s1): ReLU(inplace)\n",
      "    (drop6_s1): Dropout(p=0.5)\n",
      "  )\n",
      "  (classifier): Sequential(\n",
      "    (0): Linear(in_features=4608, out_features=4096, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Dropout(p=0.5)\n",
      "    (3): Linear(in_features=4096, out_features=9, bias=True)\n",
      "  )\n",
      "  (loss_fn): CrossEntropyLoss()\n",
      ")\n",
      "Iteration 0\n",
      "2.1721573543548582\n",
      "\t epoch= 0 \t time = 6m 16s \t lr= 0.001 \t loss =  2.1721573543548582 \t Acc = 30.806111 %\n",
      "2.163437545108795\n",
      "\t epoch= 1 \t time = 12m 34s \t lr= 0.001 \t loss =  2.163437545108795 \t Acc = 32.268056 %\n",
      "2.1564875739097595\n",
      "\t epoch= 2 \t time = 18m 54s \t lr= 0.001 \t loss =  2.1564875739097595 \t Acc = 33.834722 %\n"
     ]
    }
   ],
   "source": [
    "train_input = Dataset(train_img)\n",
    "train_loader = DataLoader(train_input, batch_size=16,shuffle=True)\n",
    "\n",
    "test_input = Dataset(test_img,loader=default_test_loader)\n",
    "test_loader = DataLoader(test_input, batch_size=batch_size,shuffle=False)\n",
    "\n",
    "model = PairPuzzle_AE(lr=init_lr,num_classes=n_crop**2)\n",
    "# model.load_state_dict(torch.load(data_dir + 'BaselinesGray/CNN_PairPuzzle_GrayAffectNet_inputsize225_minLoss.pth'))\n",
    "\n",
    "model.to(device) \n",
    "print(model)\n",
    "loss = 2.1866\n",
    "\n",
    "for i in range(1):\n",
    "    print('Iteration', i)\n",
    "    loss = model.train_model(data_loader=train_loader, num_epochs=100,lr=init_lr, min_loss=loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PairPuzzle(nn.Module):\n",
    "\n",
    "    def __init__(self,model=None,lr=0.0001,num_classes=1000):\n",
    "        super(PairPuzzle, self).__init__()\n",
    "        \n",
    "        alexnet = models.alexnet(pretrained=True)\n",
    "        \n",
    "        self.feature = alexnet.features\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(9472,4608),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4608,4096),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4096,100),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(100, num_classes)\n",
    "        )\n",
    "        \n",
    "        self.loss_fn = nn.CrossEntropyLoss() #FocalLoss(4.0)\n",
    "        self.optimizer=torch.optim.Adam([p for p in self.parameters() if p.requires_grad] , lr=lr)\n",
    "        \n",
    "\n",
    "    def forward(self,x,n_x):\n",
    "        bs, n_puzzle, c, h, w = n_x.size() \n",
    "        n_x = n_x.view(-1, c, h, w)\n",
    "        f = self.feature(n_x).view(bs,n_puzzle,-1)\n",
    "        z = self.feature(x).view(bs,1,-1)\n",
    "        z = z.expand(bs,n_puzzle,-1)\n",
    "        f_all = torch.cat((f,z),-1)\n",
    "        p = self.classifier(f_all)\n",
    "        return f_all, p.view(bs, n_puzzle, -1)\n",
    "        \n",
    "    def train_model(self, data_loader, num_epochs, lr, min_loss):\n",
    "        self.train()\n",
    "        start=time.time()\n",
    "\n",
    "        for epoch in range(num_epochs):\n",
    "            running_loss=0\n",
    "            running_acc = 0.\n",
    "            total = 0\n",
    "            num_batches = 0\n",
    "            for count, (images, n_images, label) in enumerate(data_loader):\n",
    "                self.optimizer.zero_grad()\n",
    "                \n",
    "                images=images.to(device)\n",
    "                n_images=n_images.to(device)\n",
    "                label=label.to(device)\n",
    "\n",
    "                features, scores = self.forward(images, n_images)\n",
    "\n",
    "                loss =  self.loss_fn(scores, label) \n",
    "                loss.backward()\n",
    "\n",
    "                self.optimizer.step()\n",
    "\n",
    "                running_loss += loss.detach().item()\n",
    "\n",
    "                acc = (scores.argmax(-1) == label).float().sum()\n",
    "                running_acc += acc.item()/9\n",
    "\n",
    "                total+=images.size(0)\n",
    "                num_batches += 1\n",
    "            total_loss = running_loss/num_batches\n",
    "            if total_loss<min_loss:\n",
    "                min_loss = total_loss\n",
    "                print(min_loss)\n",
    "                torch.save(model.state_dict(),data_dir+'BaselinesGray/Puzzle_GrayAffectNet_inputsize{}_minLoss.pth'.format(cut_size))\n",
    "            total_acc = running_acc/total\n",
    "            elapsed = time.time()-start\n",
    "            print('\\t', 'epoch=',epoch, '\\t time = {:.0f}m {:.0f}s'.format(elapsed // 60, elapsed % 60),'\\t lr=', lr,'\\t loss = ', total_loss , '\\t Acc = {:.6f}'.format(total_acc*100),'%')\n",
    "        return min_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_input = Dataset(train_img)\n",
    "train_loader = DataLoader(train_input, batch_size=batch_size,shuffle=True)\n",
    "\n",
    "test_input = Dataset(test_img,loader=default_test_loader)\n",
    "test_loader = DataLoader(test_input, batch_size=batch_size,shuffle=False)\n",
    "\n",
    "model = PairPuzzle(lr=init_lr,num_classes=n_crop**2)\n",
    "model.to(device)\n",
    "print(model)\n",
    "loss = 5.\n",
    "\n",
    "for i in range(1):\n",
    "    print('Iteration', i)\n",
    "    loss = model.train_model(data_loader=train_loader, num_epochs=100,lr=init_lr, min_loss=loss)"
   ]
  }
 ],
 "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}