Testing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import argparse\n",
    "import sys\n",
    "import time \n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.backends.cudnn as cudnn\n",
    "from torch.autograd import Variable\n",
    "import torch.utils.data as data\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from datasets import *\n",
    "from net.model import *\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\" \n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '3,4'\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PKU-sketech dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from config import Config_pku\n",
    "from engine.pku_market_trainer import *\n",
    "pku_args = Config_pku()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==> Loading data..\n"
     ]
    }
   ],
   "source": [
    "train_set, test_color_set, test_sketch_set = make_data_loader_pku(pku_args)\n",
    "\n",
    "n_pku_clc = len(np.unique(train_set.train_data['sketch_labels'])) # number of person in training set\n",
    "n_test_color = len(test_color_set.color_labels) #number of instance in test color set\n",
    "n_test_sketch = len(test_sketch_set.sketch_labels) #number of instance in test sketch set\n",
    "\n",
    "pku_train_loader = data.DataLoader(train_set, batch_size=pku_args.batch_size, shuffle=True, num_workers=pku_args.num_workers, drop_last=True)\n",
    "pku_test_color_loader = data.DataLoader(test_color_set, batch_size=pku_args.test_batch, shuffle=False, num_workers=pku_args.num_workers)\n",
    "pku_test_sketch_loader = data.DataLoader(test_sketch_set, batch_size=pku_args.test_batch, shuffle=False, num_workers=pku_args.num_workers)\n",
    "\n",
    "test_color_labels = test_color_set.color_labels\n",
    "test_sketch_labels = test_sketch_set.sketch_labels\n",
    "\n",
    "nquery = n_test_sketch\n",
    "ngall = n_test_color\n",
    "\n",
    "query_loader = pku_test_sketch_loader\n",
    "gall_loader = pku_test_color_loader \n",
    "query_label = test_sketch_labels\n",
    "gall_label = test_color_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "C_net(\n",
       "  (bottleneck): BatchNorm1d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (classifier): Sequential(\n",
       "    (0): Dropout(p=0.7)\n",
       "    (1): Linear(in_features=2048, out_features=150, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Embed_net_1 = Baseline(pretrain_choice=None)\n",
    "Embed_net_2 = Baseline(pretrain_choice=None)\n",
    "Classify_net = C_net(2048,n_pku_clc,0.7)\n",
    "A_net = Attribute_net(dim=2048, n_att=8)\n",
    "\n",
    "Embed_net_1.to(device)\n",
    "Embed_net_2.to(device)\n",
    "Classify_net.to(device)\n",
    "A_net.to(device)\n",
    "\n",
    "checkpoint = torch.load(pku_args.model_path+'pku_best_96.t')\n",
    "Embed_net_1.load_state_dict(checkpoint['Embed_net_1'])\n",
    "Embed_net_2.load_state_dict(checkpoint['Embed_net_2']) \n",
    "Classify_net.load_state_dict(checkpoint['Classify_net']) \n",
    "\n",
    "Embed_net_1.eval()\n",
    "Embed_net_2.eval()\n",
    "Classify_net.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting Query Feature...\n",
      "Extracting Gallery Feature...\n"
     ]
    }
   ],
   "source": [
    "print ('Extracting Query Feature...')\n",
    "ptr = 0\n",
    "query_feat = np.zeros((nquery, 2048))\n",
    "with torch.no_grad():\n",
    "    for batch_idx, (img, label ) in enumerate(query_loader):\n",
    "        batch_num = img.size(0)\n",
    "        img = img.cuda()\n",
    "        _,feat = Embed_net_2(img)\n",
    "        feat = Classify_net(feat)\n",
    "        feat = F.normalize(feat, p=2, dim=1) # normorlize features\n",
    "        query_feat[ptr:ptr+batch_num,: ] = feat.detach().cpu().numpy()\n",
    "        ptr = ptr + batch_num\n",
    "\n",
    "        \n",
    "print ('Extracting Gallery Feature...')\n",
    "ptr = 0\n",
    "gall_feat = np.zeros((ngall, 2048))\n",
    "with torch.no_grad():\n",
    "    for batch_idx, (img, label ) in enumerate(gall_loader):\n",
    "        batch_num = img.size(0)\n",
    "        img = img.cuda()\n",
    "        _,feat = Embed_net_1(img)\n",
    "        feat = Classify_net(feat)\n",
    "        feat = F.normalize(feat, p=2, dim=1) # normorlize features\n",
    "        gall_feat[ptr:ptr+batch_num,: ] = feat.detach().cpu().numpy()\n",
    "        ptr = ptr + batch_num "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CMC from 1 to 20:[0.96 0.98 0.98 0.98 1.   1.   1.   1.   1.   1.   1.   1.   1.   1.\n",
      " 1.   1.   1.   1.   1.   1.  ]\n",
      "mAP:0.9430336134453782\n"
     ]
    }
   ],
   "source": [
    "# compute the similarity\n",
    "distmat  = np.matmul(query_feat, np.transpose(gall_feat))\n",
    "cmc, mAP = eval_pku(-distmat, query_label, gall_label)\n",
    "print(f'CMC from 1 to 20:{cmc}')\n",
    "print(f'mAP:{mAP}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Shoe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from config import Config_sbir\n",
    "from engine.zap_sbir_trainer import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "sbir_args = Config_sbir()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==> Loading data..\n"
     ]
    }
   ],
   "source": [
    "train_set, test_color_set, test_sketch_set = make_data_loader_sbir(sbir_args)\n",
    "\n",
    "n_sbir_clc = len(np.unique(train_set.train_data['labels'])) \n",
    "n_test_color = len(test_color_set.color_labels) \n",
    "n_test_sketch = len(test_sketch_set.sketch_labels) \n",
    "\n",
    "\n",
    "sbir_train_loader = data.DataLoader(train_set, batch_size=sbir_args.batch_size, shuffle=True, num_workers=sbir_args.num_workers, drop_last=False)\n",
    "sbir_test_color_loader = data.DataLoader(test_color_set, batch_size=sbir_args.test_batch, shuffle=False, num_workers=sbir_args.num_workers)\n",
    "sbir_test_sketch_loader = data.DataLoader(test_sketch_set, batch_size=sbir_args.test_batch, shuffle=False, num_workers=sbir_args.num_workers)\n",
    "\n",
    "test_color_labels = np.array(test_color_set.color_labels)\n",
    "test_sketch_labels = np.array(test_sketch_set.sketch_labels)\n",
    "\n",
    "nquery = n_test_sketch\n",
    "ngall = n_test_color\n",
    "query_loader = sbir_test_sketch_loader\n",
    "gall_loader = sbir_test_color_loader \n",
    "query_label = test_sketch_labels\n",
    "gall_label = test_color_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataParallel(\n",
       "  (module): Baseline(\n",
       "    (base): SENet(\n",
       "      (layer0): Sequential(\n",
       "        (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu1): ReLU(inplace)\n",
       "        (pool): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=True)\n",
       "      )\n",
       "      (layer1): Sequential(\n",
       "        (0): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (1): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "        (2): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (layer2): Sequential(\n",
       "        (0): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (1): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "        (2): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "        (3): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (layer3): Sequential(\n",
       "        (0): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (1): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "        (2): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "        (3): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "        (4): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "        (5): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (layer4): Sequential(\n",
       "        (0): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(1024, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (1): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(2048, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "        (2): SEResNeXtBottleneck(\n",
       "          (conv1): Conv2d(2048, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv3): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (se_module): SEModule(\n",
       "            (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (relu): ReLU(inplace)\n",
       "            (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (sigmoid): Sigmoid()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (avg_pool): AvgPool2d(kernel_size=7, stride=1, padding=0)\n",
       "    )\n",
       "    (gap): AdaptiveAvgPool2d(output_size=1)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Embed_net_1 = Baseline(pretrain_choice=None)\n",
    "Embed_net_2 = Baseline(pretrain_choice=None)\n",
    "Classify_net = C_net(sbir_args.low_dim,n_sbir_clc, 0.7)\n",
    "A_net = Attribute_net(dim=sbir_args.low_dim, n_att=sbir_args.num_att)\n",
    "D_net = Domain_net(dim=sbir_args.low_dim)\n",
    "\n",
    "Embed_net_1.to(device)\n",
    "Embed_net_2.to(device)\n",
    "Classify_net.to(device)\n",
    "A_net.to(device)\n",
    "\n",
    "Embed_net_1 = nn.DataParallel(Embed_net_1)\n",
    "Embed_net_2 = nn.DataParallel(Embed_net_2)\n",
    "Classify_net = nn.DataParallel(Classify_net)\n",
    "A_net = nn.DataParallel(A_net)\n",
    "\n",
    "checkpoint = torch.load(sbir_args.model_path+'sbir_best_69.t')\n",
    "Embed_net_1.load_state_dict(checkpoint['Embed_net_1'])\n",
    "Embed_net_2.load_state_dict(checkpoint['Embed_net_2'])   \n",
    "\n",
    "Embed_net_1.eval()\n",
    "Embed_net_2.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting Query Feature...\n",
      "Extracting Gallery Feature...\n"
     ]
    }
   ],
   "source": [
    "print ('Extracting Query Feature...')\n",
    "ptr = 0\n",
    "query_feat = np.zeros((nquery, 2048))\n",
    "with torch.no_grad():\n",
    "    for batch_idx, (img, label ) in enumerate(query_loader):\n",
    "        batch_num = img.size(0)\n",
    "        img = img.cuda()\n",
    "        _,feat = Embed_net_2(img)\n",
    "        feat = F.normalize(feat, p=2, dim=1) # normorlize features\n",
    "        query_feat[ptr:ptr+batch_num,: ] = feat.detach().cpu().numpy()\n",
    "        attributes = A_net(feat)\n",
    "        ptr = ptr + batch_num         \n",
    "print ('Extracting Gallery Feature...')\n",
    "ptr = 0\n",
    "gall_feat = np.zeros((ngall, 2048))\n",
    "\n",
    "with torch.no_grad():\n",
    "    for batch_idx, (img, label ) in enumerate(gall_loader):\n",
    "        batch_num = img.size(0)\n",
    "        img = img.cuda()\n",
    "        _,feat = Embed_net_1(img)\n",
    "        feat = F.normalize(feat, p=2, dim=1) # normorlize features\n",
    "        gall_feat[ptr:ptr+batch_num,: ] = feat.detach().cpu().numpy()\n",
    "        ptr = ptr + batch_num \n",
    "\n",
    "# compute the similarity\n",
    "distmat  = np.matmul(query_feat, np.transpose(gall_feat))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CMC from 1 to 20:[0.6869565  0.79130435 0.8521739  0.87826085 0.90434784 0.93043476\n",
      " 0.9478261  0.9478261  0.9478261  0.9478261  0.95652175 0.95652175\n",
      " 0.95652175 0.9652174  0.9652174  0.9652174  0.9652174  0.9652174\n",
      " 0.9652174  0.9652174 ]\n",
      "mAP:0.7802725383916311\n"
     ]
    }
   ],
   "source": [
    "from eval_metrics import eval_sbir\n",
    "cmc, mAP = eval_sbir(-distmat, query_label, gall_label)\n",
    "print(f'CMC from 1 to 20:{cmc}')\n",
    "print(f'mAP:{mAP}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# transfer weigh from parrall to single gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "Embed_net_1 = Baseline(pretrain_choice=None)\n",
    "Embed_net_2 = Baseline(pretrain_choice=None)\n",
    "Classify_net = C_net(sbir_args.low_dim,n_sbir_clc, 0.7)\n",
    "A_net = Attribute_net(dim=sbir_args.low_dim, n_att=sbir_args.num_att)\n",
    "D_net = Domain_net(dim=sbir_args.low_dim)\n",
    "\n",
    "Embed_net_1 = nn.DataParallel(Embed_net_1)\n",
    "Embed_net_2 = nn.DataParallel(Embed_net_2)\n",
    "Classify_net = nn.DataParallel(Classify_net)\n",
    "A_net = nn.DataParallel(A_net)\n",
    "D_net = nn.DataParallel(D_net) \n",
    "\n",
    "checkpoint = torch.load(sbir_args.model_path+'zap_p.t')\n",
    "Embed_net_1.load_state_dict(checkpoint['Embed_net_1'])\n",
    "Embed_net_2.load_state_dict(checkpoint['Embed_net_2'])   \n",
    "A_net.load_state_dict(checkpoint['A_net'])\n",
    "      \n",
    "state = {\n",
    "    'Embed_net_1': Embed_net_1.module.state_dict(),\n",
    "    'Embed_net_2': Embed_net_2.module.state_dict(),\n",
    "    'Classify_net': Classify_net.module.state_dict(),\n",
    "    'A_net': A_net.module.state_dict(),\n",
    "}\n",
    "torch.save(state, 'zap.t')"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}