one_fold.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import cv2\n",
    "import glob\n",
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\"\n",
    "%pylab inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, codecs, glob\n",
    "from PIL import Image, ImageDraw\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import cv2\n",
    "\n",
    "import torch\n",
    "torch.backends.cudnn.benchmark = False\n",
    "# torch.backends.cudnn.enabled = False\n",
    "\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data.dataset import Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class XunFeiDataset(Dataset):\n",
    "    def __init__(self, img_path, img_group, transform):\n",
    "        self.img_path = img_path\n",
    "        self.transform = transform\n",
    "        self.group = img_group\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        img = Image.open(self.img_path[index]).convert('RGB')\n",
    "        \n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img)\n",
    "        \n",
    "        return img, self.group[index]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.read_csv('data/train.csv')\n",
    "train_df['path'] = 'data/train_no_face/' + train_df['name']\n",
    "\n",
    "train_df['group'] = pd.factorize(train_df['label'])[0]\n",
    "train_df['fold'] = train_df['group'] % 5\n",
    "\n",
    "train_df = train_df.sort_values(by='group')\n",
    "train_df.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr_path = train_df[train_df['fold'] != 0]['path'].values\n",
    "tr_label = train_df[train_df['fold'] != 0]['group']\n",
    "tr_label = pd.factorize(tr_label)[0]\n",
    "\n",
    "val_path = train_df[train_df['fold'] == 0]['path'].values\n",
    "val_label = train_df[train_df['fold'] == 0]['group'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr_label.max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import math\n",
    "import torchvision.transforms as T\n",
    "from torch.utils.data.sampler import Sampler\n",
    "import torch.nn.functional as F\n",
    "from collections import defaultdict\n",
    "from bisect import bisect_right\n",
    "import copy\n",
    "\n",
    "\n",
    "class RandomErasing(object):\n",
    "    \"\"\" Randomly selects a rectangle region in an image and erases its pixels.\n",
    "        'Random Erasing Data Augmentation' by Zhong et al.\n",
    "        See https://arxiv.org/pdf/1708.04896.pdf\n",
    "    Args:\n",
    "         probability: The probability that the Random Erasing operation will be performed.\n",
    "         sl: Minimum proportion of erased area against input image.\n",
    "         sh: Maximum proportion of erased area against input image.\n",
    "         r1: Minimum aspect ratio of erased area.\n",
    "         mean: Erasing value.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4914, 0.4822, 0.4465)):\n",
    "        self.probability = probability\n",
    "        self.mean = mean\n",
    "        self.sl = sl\n",
    "        self.sh = sh\n",
    "        self.r1 = r1\n",
    "\n",
    "    def __call__(self, img):\n",
    "\n",
    "        if random.uniform(0, 1) >= self.probability:\n",
    "            return img\n",
    "\n",
    "        for attempt in range(100):\n",
    "            area = img.size()[1] * img.size()[2]\n",
    "\n",
    "            target_area = random.uniform(self.sl, self.sh) * area\n",
    "            aspect_ratio = random.uniform(self.r1, 1 / self.r1)\n",
    "\n",
    "            h = int(round(math.sqrt(target_area * aspect_ratio)))\n",
    "            w = int(round(math.sqrt(target_area / aspect_ratio)))\n",
    "\n",
    "            if w < img.size()[2] and h < img.size()[1]:\n",
    "                x1 = random.randint(0, img.size()[1] - h)\n",
    "                y1 = random.randint(0, img.size()[2] - w)\n",
    "                if img.size()[0] == 3:\n",
    "                    img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]\n",
    "                    img[1, x1:x1 + h, y1:y1 + w] = self.mean[1]\n",
    "                    img[2, x1:x1 + h, y1:y1 + w] = self.mean[2]\n",
    "                else:\n",
    "                    img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]\n",
    "                return img\n",
    "\n",
    "        return img\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "SIZE_TRAIN = [320, 320]\n",
    "PROB = 0.5\n",
    "PADDING = 10\n",
    "RE_PROB = 0.2\n",
    "SIZE_TEST = [320, 320]\n",
    "\n",
    "\n",
    "def build_transforms(is_train=True):\n",
    "    normalize_transform = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "    if is_train:\n",
    "        transform = T.Compose([\n",
    "            T.Resize(SIZE_TRAIN),\n",
    "            T.RandomHorizontalFlip(PROB),\n",
    "            T.Pad(PADDING),\n",
    "            T.RandomCrop(SIZE_TRAIN),\n",
    "            T.CenterCrop(SIZE_TRAIN),\n",
    "            T.RandomRotation(5),\n",
    "            T.ToTensor(),\n",
    "            normalize_transform,\n",
    "            RandomErasing(probability=RE_PROB),\n",
    "        ])\n",
    "    else:\n",
    "        transform = T.Compose([\n",
    "            T.Resize(SIZE_TEST),\n",
    "            T.ToTensor(),\n",
    "            normalize_transform\n",
    "        ])\n",
    "\n",
    "    return transform\n",
    "\n",
    "\n",
    "# PK采样\n",
    "class RandomIdentitySampler(Sampler):\n",
    "    \"\"\"\n",
    "    Randomly sample N identities, then for each identity,\n",
    "    randomly sample K instances, therefore batch size is N*K.\n",
    "    Args:\n",
    "    - data_source (list): list of (img_path, pid, camid).\n",
    "    - num_instances (int): number of instances per identity in a batch.\n",
    "    - batch_size (int): number of examples in a batch.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, data_source, batch_size, num_instances):\n",
    "        self.data_source = data_source\n",
    "        self.batch_size = batch_size\n",
    "        self.num_instances = num_instances\n",
    "        self.num_pids_per_batch = self.batch_size // self.num_instances\n",
    "        self.index_dic = defaultdict(list)\n",
    "        for index, (_, pid) in enumerate(self.data_source):\n",
    "            self.index_dic[pid].append(index)\n",
    "        self.pids = list(self.index_dic.keys())\n",
    "\n",
    "        # estimate number of examples in an epoch\n",
    "        self.length = 0\n",
    "        for pid in self.pids:\n",
    "            idxs = self.index_dic[pid]\n",
    "            num = len(idxs)\n",
    "            if num < self.num_instances:\n",
    "                num = self.num_instances\n",
    "            self.length += num - num % self.num_instances\n",
    "\n",
    "    def __iter__(self):\n",
    "        batch_idxs_dict = defaultdict(list)\n",
    "\n",
    "        for pid in self.pids:\n",
    "            idxs = copy.deepcopy(self.index_dic[pid])\n",
    "            if len(idxs) < self.num_instances:\n",
    "                idxs = np.random.choice(idxs, size=self.num_instances, replace=True)\n",
    "            random.shuffle(idxs)\n",
    "            batch_idxs = []\n",
    "            for idx in idxs:\n",
    "                batch_idxs.append(idx)\n",
    "                if len(batch_idxs) == self.num_instances:\n",
    "                    batch_idxs_dict[pid].append(batch_idxs)\n",
    "                    batch_idxs = []\n",
    "\n",
    "        avai_pids = copy.deepcopy(self.pids)\n",
    "        final_idxs = []\n",
    "\n",
    "        while len(avai_pids) >= self.num_pids_per_batch:\n",
    "            selected_pids = random.sample(avai_pids, self.num_pids_per_batch)\n",
    "            for pid in selected_pids:\n",
    "                batch_idxs = batch_idxs_dict[pid].pop(0)\n",
    "                final_idxs.extend(batch_idxs)\n",
    "                if len(batch_idxs_dict[pid]) == 0:\n",
    "                    avai_pids.remove(pid)\n",
    "\n",
    "        self.length = len(final_idxs)\n",
    "        return iter(final_idxs)\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "    \n",
    "    \n",
    "def BaseDataset(tr_path, tr_label):\n",
    "    train_sample_data = []\n",
    "    for path, label in zip(tr_path, tr_label):\n",
    "        train_sample_data.append((path, label))\n",
    "\n",
    "    return train_sample_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = torch.utils.data.DataLoader(\n",
    "    XunFeiDataset(tr_path, tr_label, build_transforms(True)),\n",
    "    batch_size=27, sampler=RandomIdentitySampler(BaseDataset(tr_path, tr_label),batch_size=27,num_instances=3),num_workers=5,\n",
    ")\n",
    "\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    XunFeiDataset(val_path, val_label,build_transforms(False)),\n",
    "    batch_size=10, shuffle=False, num_workers=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ArcModule(nn.Module):\n",
    "    def __init__(self, in_features, out_features, s = 10, m = 0.2):\n",
    "        super().__init__()\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.s = s\n",
    "        self.m = m\n",
    "        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))\n",
    "        nn.init.xavier_normal_(self.weight)\n",
    "\n",
    "        self.cos_m = math.cos(m)\n",
    "        self.sin_m = math.sin(m)\n",
    "        self.th = torch.tensor(math.cos(math.pi - m))\n",
    "        self.mm = torch.tensor(math.sin(math.pi - m) * m)\n",
    "\n",
    "    def forward(self, inputs, labels):\n",
    "        cos_th = F.linear(inputs, F.normalize(self.weight))\n",
    "        cos_th = cos_th.clamp(-1, 1)\n",
    "        sin_th = torch.sqrt(1.0 - torch.pow(cos_th, 2))\n",
    "        cos_th_m = cos_th * self.cos_m - sin_th * self.sin_m\n",
    "        # print(type(cos_th), type(self.th), type(cos_th_m), type(self.mm))\n",
    "        cos_th_m = torch.where(cos_th > self.th, cos_th_m, cos_th - self.mm)\n",
    "        \n",
    "        cond_v = cos_th - self.th\n",
    "        cond = cond_v <= 0\n",
    "        cos_th_m[cond] = (cos_th - self.mm)[cond]\n",
    "\n",
    "        if labels.dim() == 1:\n",
    "            labels = labels.unsqueeze(-1)\n",
    "        onehot = torch.zeros(cos_th.size()).cuda()\n",
    "        labels = labels.type(torch.LongTensor).cuda()\n",
    "        onehot.scatter_(1, labels, 1.0)\n",
    "        outputs = onehot * cos_th_m + (1.0 - onehot) * cos_th\n",
    "        outputs = outputs * self.s\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import timm\n",
    "\n",
    "class XunFeiNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(XunFeiNet, self).__init__()\n",
    "                \n",
    "        model = timm.create_model('efficientnet_b3', pretrained=True)\n",
    "        model.classifier = torch.nn.Identity()\n",
    "        self.model = model\n",
    "        self.margin = ArcModule(in_features=1536, out_features = 1806)\n",
    "        \n",
    "    def forward(self, img, labels=None):        \n",
    "        global_feas = self.model(img)\n",
    "        feat = F.normalize(global_feas)\n",
    "        if labels is not None:\n",
    "            return global_feas, self.margin(feat, labels)\n",
    "        return feat\n",
    "    \n",
    "model = XunFeiNet().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义损失函数\n",
    "def normalize(x, axis=-1):\n",
    "    \"\"\"Normalizing to unit length along the specified dimension.\n",
    "    Args:\n",
    "      x: pytorch Variable\n",
    "    Returns:\n",
    "      x: pytorch Variable, same shape as input\n",
    "    \"\"\"\n",
    "    x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)\n",
    "    return x\n",
    "\n",
    "\n",
    "def euclidean_dist(x, y):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "      x: pytorch Variable, with shape [m, d]\n",
    "      y: pytorch Variable, with shape [n, d]\n",
    "    Returns:\n",
    "      dist: pytorch Variable, with shape [m, n]\n",
    "    \"\"\"\n",
    "    m, n = x.size(0), y.size(0)\n",
    "    xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)\n",
    "    yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()\n",
    "    dist = xx + yy\n",
    "    dist.addmm_(1, -2, x, y.t())\n",
    "    dist = dist.clamp(min=1e-10).sqrt()  # for numerical stability\n",
    "    return dist\n",
    "\n",
    "\n",
    "def cosine_dist(x, y):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "      x: pytorch Variable, with shape [m, d]\n",
    "      y: pytorch Variable, with shape [n, d]\n",
    "    Returns:\n",
    "      dist: pytorch Variable, with shape [m, n]\n",
    "    \"\"\"\n",
    "    m, n = x.size(0), y.size(0)\n",
    "    x_norm = torch.pow(x, 2).sum(1, keepdim=True).sqrt().expand(m, n)\n",
    "    y_norm = torch.pow(y, 2).sum(1, keepdim=True).sqrt().expand(n, m).t()\n",
    "    xy_intersection = torch.mm(x, y.t())\n",
    "    dist = xy_intersection/(x_norm * y_norm)\n",
    "    dist = (1. - dist) / 2\n",
    "    return dist\n",
    "\n",
    "\n",
    "def hard_example_mining(dist_mat, labels, return_inds=False):\n",
    "    \"\"\"For each anchor, find the hardest positive and negative sample.\n",
    "    Args:\n",
    "      dist_mat: pytorch Variable, pair wise distance between samples, shape [N, N]\n",
    "      labels: pytorch LongTensor, with shape [N]\n",
    "      return_inds: whether to return the indices. Save time if `False`(?)\n",
    "    Returns:\n",
    "      dist_ap: pytorch Variable, distance(anchor, positive); shape [N]\n",
    "      dist_an: pytorch Variable, distance(anchor, negative); shape [N]\n",
    "      p_inds: pytorch LongTensor, with shape [N];\n",
    "        indices of selected hard positive samples; 0 <= p_inds[i] <= N - 1\n",
    "      n_inds: pytorch LongTensor, with shape [N];\n",
    "        indices of selected hard negative samples; 0 <= n_inds[i] <= N - 1\n",
    "    NOTE: Only consider the case in which all labels have same num of samples,\n",
    "      thus we can cope with all anchors in parallel.\n",
    "    \"\"\"\n",
    "\n",
    "    assert len(dist_mat.size()) == 2\n",
    "    assert dist_mat.size(0) == dist_mat.size(1)\n",
    "    N = dist_mat.size(0)\n",
    "\n",
    "    # shape [N, N]\n",
    "    is_pos = labels.expand(N, N).eq(labels.expand(N, N).t())\n",
    "    is_neg = labels.expand(N, N).ne(labels.expand(N, N).t())\n",
    "\n",
    "    # `dist_ap` means distance(anchor, positive)\n",
    "    # both `dist_ap` and `relative_p_inds` with shape [N, 1]\n",
    "    dist_ap, relative_p_inds = torch.max(\n",
    "        dist_mat[is_pos].contiguous().view(N, -1), 1, keepdim=True)\n",
    "\n",
    "    # `dist_an` means distance(anchor, negative)\n",
    "    # both `dist_an` and `relative_n_inds` with shape [N, 1]\n",
    "    dist_an, relative_n_inds = torch.min(\n",
    "        dist_mat[is_neg].contiguous().view(N, -1), 1, keepdim=True)\n",
    "\n",
    "    # shape [N]\n",
    "    dist_ap = dist_ap.squeeze(1)\n",
    "    dist_an = dist_an.squeeze(1)\n",
    "\n",
    "    if return_inds:\n",
    "        # shape [N, N]\n",
    "        ind = (labels.new().resize_as_(labels)\n",
    "               .copy_(torch.arange(0, N).long())\n",
    "               .unsqueeze(0).expand(N, N))\n",
    "        # shape [N, 1]\n",
    "        p_inds = torch.gather(\n",
    "            ind[is_pos].contiguous().view(N, -1), 1, relative_p_inds.data)\n",
    "        n_inds = torch.gather(\n",
    "            ind[is_neg].contiguous().view(N, -1), 1, relative_n_inds.data)\n",
    "        # shape [N]\n",
    "        p_inds = p_inds.squeeze(1)\n",
    "        n_inds = n_inds.squeeze(1)\n",
    "        return dist_ap, dist_an, p_inds, n_inds\n",
    "\n",
    "    return dist_ap, dist_an\n",
    "\n",
    "\n",
    "class TripletLoss(object):\n",
    "    \"\"\"Modified from Tong Xiao's open-reid (https://github.com/Cysu/open-reid).\n",
    "    Related Triplet Loss theory can be found in paper 'In Defense of the Triplet\n",
    "    Loss for Person Re-Identification'.\"\"\"\n",
    "\n",
    "    def __init__(self, margin=None):\n",
    "        self.margin = margin\n",
    "        if margin is not None:\n",
    "            self.ranking_loss = nn.MarginRankingLoss(margin=margin)\n",
    "        else:\n",
    "            self.ranking_loss = nn.SoftMarginLoss()\n",
    "\n",
    "    def __call__(self, global_feat, labels, normalize_feature=False):\n",
    "        if normalize_feature:\n",
    "            global_feat = normalize(global_feat, axis=-1)\n",
    "        dist_mat = euclidean_dist(global_feat, global_feat)\n",
    "#         dist_mat = cosine_dist(global_feat, global_feat)\n",
    "\n",
    "        dist_ap, dist_an = hard_example_mining(\n",
    "            dist_mat, labels)\n",
    "        y = dist_an.new().resize_as_(dist_an).fill_(1)\n",
    "        if self.margin is not None:\n",
    "            loss = self.ranking_loss(dist_an, dist_ap, y)\n",
    "        else:\n",
    "            loss = self.ranking_loss(dist_an - dist_ap, y)\n",
    "        return loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tri_loss = TripletLoss(0.6)\n",
    "cetner_loss_weight = 0.00001\n",
    "\n",
    "\n",
    "def train(train_loader, model, criterion, optimizer,):\n",
    "    model.train()\n",
    "\n",
    "    for i, (input, target) in enumerate(train_loader):\n",
    "        input = input.cuda(non_blocking=True)\n",
    "        target = target.cuda(non_blocking=True)\n",
    "\n",
    "        gloabl_feas, output = model(input, target)\n",
    "        loss = criterion(output, target) + tri_loss(gloabl_feas, target)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        \n",
    "        if i % 40 == 0:\n",
    "            print(loss.item())\n",
    "            \n",
    "def validate(val_loader, model):\n",
    "    model.eval()\n",
    "    \n",
    "    val_feats = []\n",
    "    with torch.no_grad():\n",
    "        end = time.time()\n",
    "        for i, (input, target) in enumerate(val_loader):\n",
    "            input = input.cuda()\n",
    "            target = target.cuda()\n",
    "\n",
    "            # compute output\n",
    "            output = model(input)\n",
    "            val_feats.append(output.data.cpu().numpy())\n",
    "    return val_feats\n",
    "\n",
    "\n",
    "class CrossEntropyLabelSmooth(nn.Module):\n",
    "    \"\"\"Cross entropy loss with label smoothing regularizer.\n",
    "    Reference:\n",
    "    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.\n",
    "    Equation: y = (1 - epsilon) * y + epsilon / K.\n",
    "    Args:\n",
    "        num_classes (int): number of classes.\n",
    "        epsilon (float): weight.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, num_classes, epsilon=0.1, use_gpu=True):\n",
    "        super(CrossEntropyLabelSmooth, self).__init__()\n",
    "        self.num_classes = num_classes\n",
    "        self.epsilon = epsilon\n",
    "        self.use_gpu = use_gpu\n",
    "        self.logsoftmax = nn.LogSoftmax(dim=1)\n",
    "\n",
    "    def forward(self, inputs, targets):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)\n",
    "            targets: ground truth labels with shape (num_classes)\n",
    "        \"\"\"\n",
    "        log_probs = self.logsoftmax(inputs)\n",
    "        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1)\n",
    "        if self.use_gpu: targets = targets.cuda()\n",
    "        targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes\n",
    "        loss = (- targets * log_probs).mean(0).sum()\n",
    "        return loss\n",
    "    \n",
    "    \n",
    "class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler):\n",
    "    def __init__(\n",
    "        self,\n",
    "        optimizer,\n",
    "        milestones,\n",
    "        gamma=0.1,\n",
    "        warmup_factor=0.1,\n",
    "        warmup_iters=5,\n",
    "        warmup_method=\"linear\",\n",
    "        last_epoch=-1,\n",
    "    ):\n",
    "        if not list(milestones) == sorted(milestones):\n",
    "            raise ValueError(\n",
    "                \"Milestones should be a list of\" \" increasing integers. Got {}\",\n",
    "                milestones,\n",
    "            )\n",
    "\n",
    "        if warmup_method not in (\"constant\", \"linear\"):\n",
    "            raise ValueError(\n",
    "                \"Only 'constant' or 'linear' warmup_method accepted\"\n",
    "                \"got {}\".format(warmup_method)\n",
    "            )\n",
    "        self.milestones = milestones\n",
    "        self.gamma = gamma\n",
    "        self.warmup_factor = warmup_factor\n",
    "        self.warmup_iters = warmup_iters\n",
    "        self.warmup_method = warmup_method\n",
    "        super(WarmupMultiStepLR, self).__init__(optimizer, last_epoch)\n",
    "\n",
    "    def get_lr(self):\n",
    "        warmup_factor = 1\n",
    "        if self.last_epoch < self.warmup_iters:\n",
    "            if self.warmup_method == \"constant\":\n",
    "                warmup_factor = self.warmup_factor\n",
    "            elif self.warmup_method == \"linear\":\n",
    "                alpha = self.last_epoch / self.warmup_iters\n",
    "                warmup_factor = self.warmup_factor * (1 - alpha) + alpha\n",
    "        return [\n",
    "            base_lr\n",
    "            * warmup_factor\n",
    "            * self.gamma ** bisect_right(self.milestones, self.last_epoch)\n",
    "            for base_lr in self.base_lrs\n",
    "        ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FocalLossWithSmoothing(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            num_classes: int,\n",
    "            gamma:int = 1,\n",
    "            lb_smooth: float = 0.1,\n",
    "            size_average: bool = True,\n",
    "            ignore_index: int = None,\n",
    "            alpha: float = None):\n",
    "        \"\"\"\n",
    "        :param gamma:\n",
    "        :param lb_smooth:\n",
    "        :param ignore_index:\n",
    "        :param size_average:\n",
    "        :param alpha:\n",
    "        \"\"\"\n",
    "        super(FocalLossWithSmoothing, self).__init__()\n",
    "        self._num_classes = num_classes\n",
    "        self._gamma = gamma\n",
    "        self._lb_smooth = lb_smooth\n",
    "        self._size_average = size_average\n",
    "        self._ignore_index = ignore_index\n",
    "        self._log_softmax = nn.LogSoftmax(dim=1)\n",
    "        self._alpha = alpha\n",
    "\n",
    "        if self._num_classes <= 1:\n",
    "            raise ValueError('The number of classes must be 2 or higher')\n",
    "        if self._gamma < 0:\n",
    "            raise ValueError('Gamma must be 0 or higher')\n",
    "        if self._alpha is not None:\n",
    "            if self._alpha <= 0 or self._alpha >= 1:\n",
    "                raise ValueError('Alpha must be 0 <= alpha <= 1')\n",
    "\n",
    "    def forward(self, logits, label):\n",
    "        \"\"\"\n",
    "        :param logits: (batch_size, class, height, width)\n",
    "        :param label:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        logits = logits.float()\n",
    "        difficulty_level = self._estimate_difficulty_level(logits, label)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            label = label.clone().detach()\n",
    "            if self._ignore_index is not None:\n",
    "                ignore = label.eq(self._ignore_index)\n",
    "                label[ignore] = 0\n",
    "            lb_pos, lb_neg = 1. - self._lb_smooth, self._lb_smooth / (self._num_classes - 1)\n",
    "            lb_one_hot = torch.empty_like(logits).fill_(\n",
    "                lb_neg).scatter_(1, label.unsqueeze(1), lb_pos).detach()\n",
    "        logs = self._log_softmax(logits)\n",
    "        loss = -torch.sum(difficulty_level * logs * lb_one_hot, dim=1)\n",
    "        if self._ignore_index is not None:\n",
    "            loss[ignore] = 0\n",
    "        return loss.mean()\n",
    "\n",
    "    def _estimate_difficulty_level(self, logits, label):\n",
    "        \"\"\"\n",
    "        :param logits:\n",
    "        :param label:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        one_hot_key = torch.nn.functional.one_hot(label, num_classes=self._num_classes)\n",
    "        if len(one_hot_key.shape) == 4:\n",
    "            one_hot_key = one_hot_key.permute(0, 3, 1, 2)\n",
    "        if one_hot_key.device != logits.device:\n",
    "            one_hot_key = one_hot_key.to(logits.device)\n",
    "        pt = one_hot_key * F.softmax(logits)\n",
    "        difficulty_level = torch.pow(1 - pt, self._gamma)\n",
    "        return difficulty_level"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import normalize\n",
    "\n",
    "criterion = CrossEntropyLabelSmooth(num_classes=1806).cuda()\n",
    "\n",
    "# center = CenterLoss(num_classes=1806, feat_dim=1536)\n",
    "# optimizer_center = torch.optim.Adam(center.parameters(), lr=0.001)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), 0.0008)\n",
    "\n",
    "scheduler = WarmupMultiStepLR(optimizer=optimizer, milestones=[10, 20])\n",
    "best_acc = 0.0\n",
    "\n",
    "for epoch in range(30):\n",
    "    print('Epoch: ', epoch)\n",
    "\n",
    "    train(train_loader, model, criterion, optimizer)\n",
    "\n",
    "    scheduler.step()\n",
    "    \n",
    "    val_feats = validate(val_loader, model)\n",
    "    val_feats = np.vstack(val_feats)\n",
    "    # val_feats = normalize(val_feats)\n",
    "    \n",
    "    val_distance = []\n",
    "    for feat in val_feats:\n",
    "        dis = np.dot(feat, val_feats.T)\n",
    "        val_distance.append(dis)\n",
    "        \n",
    "    best_threahold, best_f1 = 0, 0\n",
    "    for threahold in np.linspace(0.5, 0.99, 20):\n",
    "        val_submit = []\n",
    "        for dis in val_distance[:]:\n",
    "            pred = np.where(dis > threahold)[0]\n",
    "            if len(pred) == 1:\n",
    "                ids = dis.argsort()[::-1]\n",
    "                pred = [x for x in ids[dis[ids] > 0.1]][:2]\n",
    "\n",
    "            val_submit.append(pred)\n",
    "\n",
    "        val_f1s = []\n",
    "        for x, pred in zip(val_label, val_submit):\n",
    "            label = np.where(val_label == x)[0]\n",
    "            val_f1 = len(set(pred) & set(label)) / len(set(pred) | set(label)) \n",
    "            val_f1s.append(val_f1)\n",
    "\n",
    "        if best_f1 < np.mean(val_f1s):\n",
    "            best_f1 = np.mean(val_f1s)\n",
    "            best_threahold = threahold\n",
    "\n",
    "    print('Val', best_threahold, best_f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "torch.save(model.state_dict(),'one_final.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}