pr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d6fcf06f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\"\n",
    "from logging import log\n",
    "import re\n",
    "import numpy as np\n",
    "import os.path as op\n",
    "from hashing_module.triplet_loss import *\n",
    "from torch.autograd.grad_mode import F\n",
    "\n",
    "from torch.nn.modules import loss\n",
    "from torch.utils.data.sampler import Sampler\n",
    "\n",
    "import argparse\n",
    "from oscar.modeling.modeling_bert import HashingformerALL,normal_label\n",
    "from pytorch_transformers import BertTokenizer, BertConfig\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler\n",
    "from tqdm import tqdm\n",
    "from torch.autograd import Variable\n",
    "from oscar.utils.tsv_file import TSVFile\n",
    "from torch.nn import CrossEntropyLoss\n",
    "import json\n",
    "import base64\n",
    "import random\n",
    "from pytorch_transformers import AdamW, WarmupLinearSchedule, WarmupConstantSchedule\n",
    "from hashing_module.utils import calc_map_k\n",
    "from oscar.utils.logger import setup_logger\n",
    "from oscar.utils.misc import mkdir\n",
    "from torch.nn import functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e15b4c95",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Opt():\n",
    "    def __init__(self) -> None:\n",
    "        self.use_gpu = True\n",
    "        self.training_size = 10000\n",
    "        self.query_size = 2000\n",
    "        self.bit = 64\n",
    "        self.database_size = 18000 \n",
    "        self.gamma = 1\n",
    "        self.eta = 1\n",
    "        self.valid = True\n",
    "        self.batch_size = 64\n",
    "        self.margin = 0.4\n",
    "        self.gamma = 1\n",
    "        self.beta = 1\n",
    "        self.alpha = 1\n",
    "opt = Opt()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "69904fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_random_seed(seed):\n",
    "    \"\"\"Set random seed.\n",
    "\n",
    "    Args:\n",
    "        seed (int): Seed to be used.\n",
    "        deterministic (bool): Whether to set the deterministic option for\n",
    "            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`\n",
    "            to True and `torch.backends.cudnn.benchmark` to False.\n",
    "            Default: False.\n",
    "    \"\"\"\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "set_random_seed(1000)\n",
    "class RetrievalDataset(Dataset):\n",
    "    \"\"\" Image/Text Retrieval Dataset\"\"\"\n",
    "    def __init__(self, args,tokenizer,split=\"train\"):\n",
    "        \"\"\"\n",
    "        tokenizer: tokenizer to process caption text.\n",
    "        args: configureation parameters including max_seq_length, etc.\n",
    "        split: used to infer the data used for training or testing. \n",
    "             All files are in .pt format of a dictionary with image keys and \n",
    "             image features (pytorch tensors), captions (list of str, support multiple\n",
    "             captions per image), labels (list of dictionary or str of all labels),\n",
    "\n",
    "        \"\"\"\n",
    "        super(RetrievalDataset, self).__init__()\n",
    "        with open(args.tagslabel ,\"r\") as f:\n",
    "            self.tagslabel = json.load(f)   \n",
    "        self.args = args  \n",
    "        self.split = split      \n",
    "        self.img_file = args.img_feat_file\n",
    "        self.img_tsv = TSVFile(self.img_file)\n",
    "        self.img_keys = list(self.tagslabel.keys())  # img_id as int\n",
    "        imgid2idx_file = op.join(op.dirname(self.img_file), 'imageid2idx.json')\n",
    "        self.image_id2idx = json.load(open(imgid2idx_file))  # img_id as string\n",
    "        with open(args.class_name,\"r\") as f:\n",
    "            self.class_name = json.load(f)\n",
    "            self.class_name = np.array(self.class_name)\n",
    "        if(args.split_keys):\n",
    "            with open(args.split_keys,\"r\") as f:\n",
    "                self.img_keys = json.load(f)\n",
    "                self.img_keys = [str(i) for i in self.img_keys]\n",
    "        else:\n",
    "            random.seed(279834)\n",
    "            random.shuffle(self.img_keys)\n",
    "        if(split==\"train\"):\n",
    "            self.img_keys = self.img_keys[args.query_size:args.training_size + args.query_size]\n",
    "        elif(split==\"query\"):\n",
    "            self.img_keys = self.img_keys[:args.query_size]\n",
    "        else:\n",
    "            self.img_keys = self.img_keys[args.query_size:args.database_size+ args.query_size]\n",
    "        label_data_dir = op.dirname(self.img_file)\n",
    "        label_file = os.path.join(label_data_dir, \"label.tsv\")\n",
    "        self.label_tsv = TSVFile(label_file)\n",
    "        self.labels = {}\n",
    "        for line_no in tqdm(range(self.label_tsv.num_rows())):\n",
    "            row = self.label_tsv.seek(line_no)\n",
    "            image_id = row[0]\n",
    "            if image_id in self.img_keys:\n",
    "                results = json.loads(row[1])\n",
    "                objects = results['objects'] if type(\n",
    "                    results) == dict else results\n",
    "                self.labels[image_id] = {\n",
    "                    \"image_h\": results[\"image_h\"] if type(\n",
    "                        results) == dict else 600,\n",
    "                    \"image_w\": results[\"image_w\"] if type(\n",
    "                        results) == dict else 800,\n",
    "                    \"class\": [cur_d['class'] for cur_d in objects],\n",
    "                    \"boxes\": np.array([cur_d['rect'] for cur_d in objects],\n",
    "                                        dtype=np.float32)\n",
    "                }\n",
    "        self.label_tsv._fp.close()\n",
    "        self.label_tsv._fp = None   \n",
    "        self.output_mode = 'classification'\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_seq_length = 35\n",
    "        self.max_img_seq_len = 70\n",
    "        self.args.max_label_length = args.max_label_length\n",
    "    def get_od_labels(self, img_key):\n",
    "\n",
    "        if type(self.labels[img_key]) == str:\n",
    "            od_labels = self.labels[img_key]\n",
    "        else:\n",
    "            od_labels = ' '.join(self.labels[img_key]['class'])\n",
    "        return od_labels\n",
    "    def class_tokenize(self,labels,max_length=15):\n",
    "        all_size = labels.shape[0]\n",
    "        final_label = []\n",
    "        for i in range(all_size):\n",
    "            this_label = torch.zeros((max_length+2))\n",
    "            class_name = self.class_name[labels[i]>0]\n",
    "            \n",
    "            tokens = self.tokenizer.tokenize(\"\".join(class_name))\n",
    "            tokens = [self.tokenizer.cls_token] + tokens[0:max_length] + [self.tokenizer.sep_token]\n",
    "            input_ids = self.tokenizer.convert_tokens_to_ids(tokens)\n",
    "            \n",
    "            this_label[0:len(input_ids)] = torch.Tensor(input_ids)\n",
    "            final_label.append(this_label)\n",
    "        final_label = torch.stack(final_label).long()  \n",
    "        return final_label  \n",
    "\n",
    "    # def class_tokenize_list(self,labels,max_length=10):\n",
    "    #     this_label = torch.zeros((max_length+2))\n",
    "    #     class_name = self.class_name[labels>0]\n",
    "    #     tokens = [self.tokenizer.cls_token] + class_name.tolist()[0:max_length] + [self.tokenizer.sep_token]\n",
    "    #     input_ids = self.tokenizer.convert_tokens_to_ids(tokens)\n",
    "    #     this_label[0:len(input_ids)] = torch.Tensor(input_ids)\n",
    "\n",
    "    #     return this_label    \n",
    "    def tensorize_example(self, text_a, img_feat, text_b=None, \n",
    "            cls_token_segment_id=0, pad_token_segment_id=0,\n",
    "            sequence_a_segment_id=0, sequence_b_segment_id=1):\n",
    "        tokens_a = self.tokenizer.tokenize(text_a)\n",
    "        if len(tokens_a) > self.max_seq_length - 2:\n",
    "            tokens_a = tokens_a[:(self.max_seq_length - 2)]\n",
    "\n",
    "        tokens = [self.tokenizer.cls_token] + tokens_a + [self.tokenizer.sep_token]\n",
    "        segment_ids = [cls_token_segment_id] + [sequence_a_segment_id] * (len(tokens_a) + 1)\n",
    "    \n",
    "        if text_b:\n",
    "            tokens_b = self.tokenizer.tokenize(text_b)\n",
    "            if len(tokens_b) > self.max_seq_length   - 2:#a\n",
    "                tokens_b = tokens_b[: (self.max_seq_length  - 2)]\n",
    "            tokens_b = [self.tokenizer.cls_token] +tokens_b+ [self.tokenizer.sep_token]\n",
    "            segment_ids_b = [sequence_b_segment_id] + [sequence_b_segment_id] * (len(tokens_b) -1)\n",
    "        #这儿分a padding\n",
    "        seq_len_a = len(tokens)\n",
    "        seq_padding_len_a = self.max_seq_length - seq_len_a\n",
    "        tokens += [self.tokenizer.pad_token] * seq_padding_len_a\n",
    "        segment_ids += [pad_token_segment_id] * seq_padding_len_a\n",
    "        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)\n",
    "        #b padding\n",
    "        seq_len_b = len(tokens_b)\n",
    "        seq_padding_len_b = self.max_seq_length - seq_len_b\n",
    "        tokens_b += [self.tokenizer.pad_token] * seq_padding_len_b\n",
    "        segment_ids_b += [pad_token_segment_id] * seq_padding_len_b\n",
    "        input_ids_b = self.tokenizer.convert_tokens_to_ids(tokens_b)\n",
    "        #合并\n",
    "        input_ids = input_ids+input_ids_b\n",
    "        segment_ids = segment_ids+segment_ids_b\n",
    "        # image features\n",
    "        img_len = img_feat.shape[0]\n",
    "\n",
    "        if img_len > self.max_img_seq_len:\n",
    "            img_feat = img_feat[0 : self.max_img_seq_len, :]\n",
    "            img_len = img_feat.shape[0]\n",
    "            img_padding_len = 0\n",
    "        else:\n",
    "            img_padding_len = self.max_img_seq_len - img_len\n",
    "            padding_matrix = torch.zeros((img_padding_len, img_feat.shape[1]))\n",
    "            \n",
    "            img_feat = torch.cat((img_feat, padding_matrix), 0)\n",
    "\n",
    "        # generate attention_mask\n",
    "        att_mask_type = \"CLR\"\n",
    "        if att_mask_type == \"CLR\":\n",
    "            attention_mask = [1] * seq_len_a + [0] * seq_padding_len_a +[1] * seq_len_b + [0] * seq_padding_len_b +  [1] * img_len + [0] * img_padding_len \n",
    "\n",
    "\n",
    "        input_ids = torch.tensor(input_ids, dtype=torch.long)\n",
    "        attention_mask = torch.tensor(attention_mask, dtype=torch.long)\n",
    "        segment_ids = torch.tensor(segment_ids, dtype=torch.long)\n",
    "        return (input_ids, attention_mask, segment_ids, img_feat)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \n",
    "        img_key = self.img_keys[index]\n",
    "      \n",
    "        feature = self.get_image(img_key)\n",
    "        tag_list = self.tagslabel[img_key][\"tags\"]\n",
    "        if(isinstance(tag_list,list)):\n",
    "            caption = \"\"\n",
    "            for i in tag_list:\n",
    "                caption+=i+\" \"\n",
    "            caption=caption.strip()\n",
    "        else:#is a string\n",
    "            caption  = tag_list\n",
    "        od_labels = self.get_od_labels(img_key)\n",
    "        example = self.tensorize_example(caption, feature, text_b=od_labels)\n",
    "        label = self.tagslabel[img_key][\"label\"]\n",
    "        label=torch.tensor(label, dtype=torch.long)\n",
    "        if(self.split==\"train\"):\n",
    "            raw_label = self.generate_samples(label,self.args.negative_number)\n",
    "            #negative_label = normal_label(raw_label,max_length = self.args.max_label_length )\n",
    "            negative_label = self.class_tokenize(raw_label,max_length = self.args.max_label_length )\n",
    "            return tuple(list(example) + [label,negative_label,raw_label]),img_key\n",
    "        else:\n",
    "            return tuple(list(example) + [label]),img_key\n",
    "\n",
    "    def generate_samples(self, label,negative_number = 99):\n",
    "        mask = 1-label #\n",
    "        negative_samples = []\n",
    "        positive_sample = label \n",
    "        while(len(negative_samples)<negative_number):\n",
    "            smaples=  torch.from_numpy(np.random.choice(2, self.args.class_number,p=[1-5/self.args.class_number,5/self.args.class_number]))\n",
    "            is_positive = (smaples*label).sum()>0\n",
    "            if(not is_positive):\n",
    "                negative_samples.append(smaples)\n",
    "        final_sample = [positive_sample]+negative_samples\n",
    "        final_sample = torch.stack(final_sample)\n",
    "        return final_sample\n",
    "    def get_image(self, image_id):\n",
    "        image_idx = self.image_id2idx[str(image_id)]\n",
    "        row = self.img_tsv.seek(image_idx)\n",
    "        num_boxes = int(row[1])\n",
    "        features = np.frombuffer(base64.b64decode(row[-1]),\n",
    "                                 dtype=np.float32).reshape((num_boxes, -1)).copy()\n",
    "        t_features = torch.from_numpy(features)\n",
    "        return t_features\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_keys) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f09705c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_neighbor(label1, label2):\n",
    "    # calculate the similar matrix\n",
    "    label1=label1.to(torch.float32)\n",
    "    label2=label2.to(torch.float32)\n",
    "    Sim = (label1.matmul(label2.transpose(0, 1)) > 0)\n",
    "    return Sim\n",
    "def save_checkpoint(model, tokenizer, args, epoch):\n",
    "    checkpoint_dir = op.join(args.output_dir, 'checkpoint-{}'.format(\n",
    "        epoch))\n",
    "    mkdir(checkpoint_dir)\n",
    "    #model_to_save = model.module if hasattr(model, 'module') else model\n",
    "    save_num = 0\n",
    "    while (save_num < 10):\n",
    "        try:\n",
    "            \n",
    "            save_pretrained(model=model,save_directory=checkpoint_dir)\n",
    "            torch.save(args, op.join(checkpoint_dir, 'training_args.bin'))\n",
    "            tokenizer.save_pretrained(checkpoint_dir)\n",
    "            logger.info(\"Save checkpoint to {}\".format(checkpoint_dir))\n",
    "            break\n",
    "        except:\n",
    "            save_num += 1\n",
    "    if save_num == 10:\n",
    "        logger.info(\"Failed to save checkpoint after 10 trails.\")\n",
    "    return\n",
    "def log_loss_func(a,b,a_l,b_l):\n",
    "    logit_it = torch.matmul(a,b.t())\n",
    "    sim_it = torch.matmul(a_l,b_l.t())>0\n",
    "    theta_it = 1/2*logit_it\n",
    "    loss = -torch.mean((sim_it * theta_it - torch.log(1.0 + torch.exp(theta_it))))\n",
    "    precision = f1_calc(logit_it,sim_it,0)\n",
    "    return loss,precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9f2e4222",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_code(model, query_dataloader,TorI=\"I\"):\n",
    "\n",
    "    class_logits = []\n",
    "    labels = []\n",
    "    hashing_bit = []\n",
    "    for batch,keys in tqdm(query_dataloader):\n",
    "        #image = X[ind]#.unsqueeze(1).unsqueeze(-1).type(torch.float)\n",
    "        train_input_ids_this = batch[0].long().cuda()\n",
    "        train_attention_mask_this = batch[1].long().cuda()\n",
    "        train_token_type_ids_this = batch[2].long().cuda()\n",
    "        train_img_feats_this = batch[3].cuda()\n",
    "        label = batch[4].cuda()\n",
    "        with torch.no_grad():\n",
    "            if(TorI==\"T\"):\n",
    "                cur_f= model(input_ids=train_input_ids_this,token_type_ids=train_token_type_ids_this,\n",
    "                                        attention_mask=train_attention_mask_this,img_feats=train_img_feats_this,modal=\"t\")\n",
    "            else:\n",
    "                cur_f= model(input_ids=train_input_ids_this,token_type_ids=train_token_type_ids_this,\n",
    "                                        attention_mask=train_attention_mask_this,img_feats=train_img_feats_this,modal=\"i\")\n",
    "            hashing_bit.append(cur_f)\n",
    "            labels.append(label)\n",
    "    hashing_bit = torch.cat(hashing_bit,0)\n",
    "    labels = torch.cat(labels,0)\n",
    "\n",
    "    #B = torch.sign(B)\n",
    "    return hashing_bit,labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5c6e9dfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\"--num_train_epochs\", default=1000, type=int, \n",
    "                    help=\"Total number of training epochs to perform.\")\n",
    "parser.add_argument(\"--class_name\", default='/raid/data_modal/MIR_Flickr_25k/class_name.json\"', type=str, required=False,\n",
    "                    help=\"The input data dir with all required files.\")\n",
    "parser.add_argument(\"--split_keys\", default='', type=str, required=False,\n",
    "                    help=\"split_keys\")\n",
    "parser.add_argument(\"--tagslabel\", default='MIR_Flickr_25k/img_tagslabel.json', type=str, required=False,\n",
    "                    help=\"The input data dir with all required files.\")\n",
    "parser.add_argument(\"--img_feat_file\", default='/MIR_Flickr_25k/vinvl_data/vinvl_vg_x152c4/predictions.tsv', type=str, required=False,\n",
    "                    help=\"The absolute address of the image feature file.\")\n",
    "parser.add_argument(\"--output_dir\", default='output/log_aipr', type=str, required=False,\n",
    "                    help=\"The output directory to save checkpoint and test results.\")\n",
    "parser.add_argument(\"--num_workers\", default=4, type=int, help=\"Workers in dataloader.\")\n",
    "parser.add_argument(\"--eval_model_dir\", type=str, default='', \n",
    "                    help=\"Model directory for evaluation.\")   \n",
    "parser.add_argument(\"--do_lower_case\", action='store_true', \n",
    "                    help=\"Set this flag if you are using an uncased model.\")   \n",
    "parser.add_argument(\"--output_file\", type=str, default='', \n",
    "                    help=\"Model directory for evaluation.\")  \n",
    "parser.add_argument(\"--learning_rate\", default=2e-5, type=float, help=\"The initial lr.\")\n",
    "parser.add_argument(\"--weight_decay\", default=0.05, type=float, help=\"Weight deay.\")\n",
    "parser.add_argument(\"--adam_epsilon\", default=1e-8, type=float, help=\"Epsilon for Adam.\")\n",
    "parser.add_argument(\"--max_grad_norm\", default=1.0, type=float, help=\"Max gradient norm.\")\n",
    "parser.add_argument(\"--warmup_steps\", default=0, type=int, help=\"Linear warmup.\")\n",
    "parser.add_argument(\"--scheduler\", default='linear', type=str, help=\"constant or linear.\")\n",
    "parser.add_argument(\"--bit\", default=64, type=int, help=\"constant or linear.\")\n",
    "parser.add_argument(\"--class_number\", default=255, type=int, help=\"constant or linear.\")  \n",
    "parser.add_argument(\"--training_size\", default=10000, type=int, help=\"constant or linear.\") \n",
    "parser.add_argument(\"--query_size\", default=2000, type=int, help=\"constant or linear.\") \n",
    "parser.add_argument(\"--database_size\", default=18000, type=int, help=\"constant or linear.\")\n",
    "parser.add_argument(\"--no_pretrain\", action='store_true', help=\"constant or linear.\")\n",
    "parser.add_argument(\"--negative_number\", default=9, type=int, help=\"constant or linear.\")\n",
    "parser.add_argument(\"--max_label_length\", default=10, type=int, help=\"constant or linear.\")\n",
    "# lists =[\n",
    "#                 \"--tagslabel\",\n",
    "#                 \"/raid/data_modal/NUS-WIDE/hashing/img_tagslabel.json\",\n",
    "#                 \"--img_feat_file\",\n",
    "#                 \"/raid/data_modal/NUS-WIDE/hashing/use_for_vinvl/feature.tsv\",\n",
    "#                 \"--do_lower_case\",\n",
    "#                 \"--output_dir\",\n",
    "#                 \"output/try_contrastive_COCO\",\n",
    "#                 \"--eval_model_dir\",\n",
    "#                 \"/raid/whf/ObjectContrastiveTransformer/output/try_contrastive_NUS_16/checkpoint-10\",\n",
    "#                 \"--training_size\",\"10500\",\n",
    "#                 \"--query_size\",\"2100\",\n",
    "#                 \"--database_size\",\"188321\",\n",
    "#                 \"--class_number\",\"21\",\n",
    "#                 \"--class_name\",\"/raid/data_modal/NUS-WIDE/hashing/class_name.json\",\n",
    "#                 \"--bit\",\"16\"\n",
    "                \n",
    "#             ]\n",
    "\n",
    "# args = parser.parse_args(lists)\n",
    "# args.output = \"output/try_contrastive_NUS_16\"\n",
    "\n",
    "\n",
    "lists =[\n",
    "                \"--tagslabel\",\n",
    "                \"/raid/data_modal/IAPR_TC-12/img_tagslabel.json\",\n",
    "                \"--img_feat_file\",\n",
    "                \"/raid/data_modal/IAPR_TC-12/vinvl/use_for_vinvl/feature.tsv\",\n",
    "                \"--do_lower_case\",\n",
    "                \"--output_dir\",\n",
    "                \"output/try_contrastive_IAPR\",\n",
    "                \"--eval_model_dir\",\n",
    "                \"/raid/whf/ObjectContrastiveTransformer/output/IAPR_16/checkpoint-150\",\n",
    "                \"--training_size\",\"10000\",\n",
    "                \"--query_size\",\"2000\",\n",
    "                \"--database_size\",\"18000\",\n",
    "                \"--class_number\",\"255\",\n",
    "                \"--class_name\",\"/raid/data_modal/IAPR_TC-12/class_name.json\",\n",
    "                \"--bit\",\"16\"\n",
    "            ]\n",
    "\n",
    "args = parser.parse_args(lists)\n",
    "args.output = \"output/IAPR_16\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b392fafa",
   "metadata": {},
   "outputs": [],
   "source": [
    "if(args.training_size != -1):\n",
    "    opt.training_size = args.training_size\n",
    "    opt.query_size = args.query_size\n",
    "    opt.database_size = args.database_size\n",
    "global logger\n",
    "mkdir(args.output_dir)\n",
    "logger = setup_logger(\"vlpretrain\", args.output_dir, 0)\n",
    "opt.bit = args.bit\n",
    "opt.class_number = args.class_number\n",
    "device = torch.device(\"cuda\")\n",
    "config_class, tokenizer_class = BertConfig, BertTokenizer\n",
    "checkpoint = args.eval_model_dir\n",
    "tokenizer = tokenizer_class.from_pretrained(checkpoint, do_lower_case=args.do_lower_case)\n",
    "config = config_class.from_pretrained(checkpoint)\n",
    "config.class_number = opt.class_number\n",
    "config.bit = args.bit\n",
    "model = HashingformerALL(None,config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "821b58f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 17585/17585 [00:01<00:00, 11397.60it/s]\n",
      "100%|██████████| 17585/17585 [00:07<00:00, 2456.70it/s]\n"
     ]
    }
   ],
   "source": [
    "if(not args.no_pretrain):\n",
    "\n",
    "    if(not os.path.exists(checkpoint+\"/pytorch_model.bin\")):\n",
    "        sd = torch.load(checkpoint+\"/model.cpkt\", map_location=\"cpu\")\n",
    "    else:\n",
    "        sd = torch.load(checkpoint+\"/pytorch_model.bin\", map_location=\"cpu\")\n",
    "    missing, unexpected = model.load_state_dict(sd, strict=False)  \n",
    "model.to(device)\n",
    "query_dataset = RetrievalDataset(args,tokenizer,\"query\")\n",
    "retrieval_dataset = RetrievalDataset(args,tokenizer,\"retrieval\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "588a85a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [00:05<00:00,  1.47s/it]\n",
      "100%|██████████| 4/4 [00:03<00:00,  1.02it/s]\n",
      "100%|██████████| 31/31 [00:33<00:00,  1.08s/it]\n",
      "100%|██████████| 31/31 [00:17<00:00,  1.72it/s]\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "query_sampler = SequentialSampler(query_dataset)\n",
    "query_dataloader = DataLoader(query_dataset, sampler=query_sampler,\n",
    "        batch_size=512, num_workers=4)\n",
    "retrieval_sampler = SequentialSampler(retrieval_dataset)\n",
    "retrieval_dataloader = DataLoader(retrieval_dataset, sampler=retrieval_sampler,\n",
    "        batch_size=512, num_workers=4)\n",
    "qBX,query_L_i = generate_code(model,query_dataloader ,TorI=\"I\")\n",
    "qBY,query_L_t = generate_code(model, query_dataloader, TorI=\"T\")\n",
    "rBX,retrieval_L_i = generate_code(model, retrieval_dataloader,TorI=\"I\")\n",
    "rBY,retrieval_L_t = generate_code(model, retrieval_dataloader, TorI=\"T\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d2b8c4ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hashing_module.utils import calc_map_k_final"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "9d1dfbaf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 12/2000 [00:00<00:17, 113.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "calc map k\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2000/2000 [00:17<00:00, 115.00it/s]\n",
      "  1%|          | 12/2000 [00:00<00:16, 117.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "calc map k\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2000/2000 [00:17<00:00, 115.45it/s]\n"
     ]
    }
   ],
   "source": [
    "mapi2t,_,pr_i2t = calc_map_k_final(torch.sign(qBX), torch.sign(rBY), query_L_i, retrieval_L_t)\n",
    "mapt2i,_,pr_t2i = calc_map_k_final(torch.sign(qBY), torch.sign(rBX), query_L_t, retrieval_L_i)\n",
    "pr={\"pr_i2t\":pr_i2t,\"pr_t2i\":pr_t2i}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "377f8a00",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(args.output,\"pr.json\"),\"w\") as f:\n",
    "    json.dump(pr,f)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "a83f8ef4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'pr_i2t': [{'TP': 613,\n",
       "   'FP': 55,\n",
       "   'TN': 21717461,\n",
       "   'FN': 9451871,\n",
       "   'P': 0.9176646706586826,\n",
       "   'R': 6.485067840368733e-05},\n",
       "  {'TP': 9023,\n",
       "   'FP': 781,\n",
       "   'TN': 21716735,\n",
       "   'FN': 9443461,\n",
       "   'P': 0.9203386372909017,\n",
       "   'R': 0.0009545639008751562},\n",
       "  {'TP': 79675,\n",
       "   'FP': 8760,\n",
       "   'TN': 21708756,\n",
       "   'FN': 9372809,\n",
       "   'P': 0.900944196302369,\n",
       "   'R': 0.008429001308015967},\n",
       "  {'TP': 339599,\n",
       "   'FP': 51137,\n",
       "   'TN': 21666379,\n",
       "   'FN': 9112885,\n",
       "   'P': 0.8691264690225625,\n",
       "   'R': 0.0359269584587501},\n",
       "  {'TP': 966267,\n",
       "   'FP': 183098,\n",
       "   'TN': 21534418,\n",
       "   'FN': 8486217,\n",
       "   'P': 0.8406963845253683,\n",
       "   'R': 0.1022236059854743},\n",
       "  {'TP': 1968924,\n",
       "   'FP': 445459,\n",
       "   'TN': 21272057,\n",
       "   'FN': 7483560,\n",
       "   'P': 0.8154977897044504,\n",
       "   'R': 0.2082969936791218},\n",
       "  {'TP': 3253019,\n",
       "   'FP': 930811,\n",
       "   'TN': 20786705,\n",
       "   'FN': 6199465,\n",
       "   'P': 0.7775217922334321,\n",
       "   'R': 0.3441443540131885},\n",
       "  {'TP': 4662885,\n",
       "   'FP': 1936236,\n",
       "   'TN': 19781280,\n",
       "   'FN': 4789599,\n",
       "   'P': 0.7065918324576864,\n",
       "   'R': 0.4932973174035523},\n",
       "  {'TP': 6182906,\n",
       "   'FP': 4127318,\n",
       "   'TN': 17590198,\n",
       "   'FN': 3269578,\n",
       "   'P': 0.5996868739224288,\n",
       "   'R': 0.6541038313315315},\n",
       "  {'TP': 7639626,\n",
       "   'FP': 8046091,\n",
       "   'TN': 13671425,\n",
       "   'FN': 1812858,\n",
       "   'P': 0.48704346763364403,\n",
       "   'R': 0.8082135870317263},\n",
       "  {'TP': 8687683,\n",
       "   'FP': 13179497,\n",
       "   'TN': 8538019,\n",
       "   'FN': 764801,\n",
       "   'P': 0.39729324951822775,\n",
       "   'R': 0.9190899450345539},\n",
       "  {'TP': 9216527,\n",
       "   'FP': 17746264,\n",
       "   'TN': 3971252,\n",
       "   'FN': 235957,\n",
       "   'P': 0.3418239231984552,\n",
       "   'R': 0.9750375668448632},\n",
       "  {'TP': 9399605,\n",
       "   'FP': 20408314,\n",
       "   'TN': 1309202,\n",
       "   'FN': 52879,\n",
       "   'P': 0.3153391888913815,\n",
       "   'R': 0.9944058090973759},\n",
       "  {'TP': 9445223,\n",
       "   'FP': 21451539,\n",
       "   'TN': 265977,\n",
       "   'FN': 7261,\n",
       "   'P': 0.30570268172438264,\n",
       "   'R': 0.9992318421274239},\n",
       "  {'TP': 9451781,\n",
       "   'FP': 21684839,\n",
       "   'TN': 32677,\n",
       "   'FN': 703,\n",
       "   'P': 0.30355835026409417,\n",
       "   'R': 0.999925628014816},\n",
       "  {'TP': 9452468,\n",
       "   'FP': 21715307,\n",
       "   'TN': 2209,\n",
       "   'FN': 16,\n",
       "   'P': 0.3032769583327652,\n",
       "   'R': 0.999998307323239}],\n",
       " 'pr_t2i': [{'TP': 431,\n",
       "   'FP': 46,\n",
       "   'TN': 21717470,\n",
       "   'FN': 9452053,\n",
       "   'P': 0.9035639412997903,\n",
       "   'R': 4.559648024794329e-05},\n",
       "  {'TP': 6420,\n",
       "   'FP': 988,\n",
       "   'TN': 21716528,\n",
       "   'FN': 9446064,\n",
       "   'P': 0.8666306695464363,\n",
       "   'R': 0.0006791865503289929},\n",
       "  {'TP': 66191,\n",
       "   'FP': 11554,\n",
       "   'TN': 21705962,\n",
       "   'FN': 9386293,\n",
       "   'P': 0.8513859412180848,\n",
       "   'R': 0.007002497967729964},\n",
       "  {'TP': 293551,\n",
       "   'FP': 71078,\n",
       "   'TN': 21646438,\n",
       "   'FN': 9158933,\n",
       "   'P': 0.8050676166733858,\n",
       "   'R': 0.031055434740751744},\n",
       "  {'TP': 836394,\n",
       "   'FP': 253286,\n",
       "   'TN': 21464230,\n",
       "   'FN': 8616090,\n",
       "   'P': 0.7675592834593642,\n",
       "   'R': 0.08848404292458999},\n",
       "  {'TP': 1705312,\n",
       "   'FP': 603840,\n",
       "   'TN': 21113676,\n",
       "   'FN': 7747172,\n",
       "   'P': 0.7385014065769598,\n",
       "   'R': 0.18040887453498997},\n",
       "  {'TP': 2869419,\n",
       "   'FP': 1207737,\n",
       "   'TN': 20509779,\n",
       "   'FN': 6583065,\n",
       "   'P': 0.7037795463308247,\n",
       "   'R': 0.30356242866954336},\n",
       "  {'TP': 4224931,\n",
       "   'FP': 2333812,\n",
       "   'TN': 19383704,\n",
       "   'FN': 5227553,\n",
       "   'P': 0.6441677925175602,\n",
       "   'R': 0.4469651575183835},\n",
       "  {'TP': 5777722,\n",
       "   'FP': 4573116,\n",
       "   'TN': 17144400,\n",
       "   'FN': 3674762,\n",
       "   'P': 0.5581888152437513,\n",
       "   'R': 0.6112384850373722},\n",
       "  {'TP': 7346613,\n",
       "   'FP': 8419921,\n",
       "   'TN': 13297595,\n",
       "   'FN': 2105871,\n",
       "   'P': 0.46596246200972263,\n",
       "   'R': 0.7772150685470612},\n",
       "  {'TP': 8529334,\n",
       "   'FP': 13383857,\n",
       "   'TN': 8333659,\n",
       "   'FN': 923150,\n",
       "   'P': 0.3892328597875134,\n",
       "   'R': 0.9023378405083785},\n",
       "  {'TP': 9159510,\n",
       "   'FP': 17783265,\n",
       "   'TN': 3934251,\n",
       "   'FN': 292974,\n",
       "   'P': 0.3399616409222881,\n",
       "   'R': 0.9690056074149398},\n",
       "  {'TP': 9386349,\n",
       "   'FP': 20396536,\n",
       "   'TN': 1320980,\n",
       "   'FN': 66135,\n",
       "   'P': 0.3151591593628354,\n",
       "   'R': 0.9930034264009333},\n",
       "  {'TP': 9442701,\n",
       "   'FP': 21435326,\n",
       "   'TN': 282190,\n",
       "   'FN': 9783,\n",
       "   'P': 0.30580648821895257,\n",
       "   'R': 0.99896503395298},\n",
       "  {'TP': 9451848,\n",
       "   'FP': 21676922,\n",
       "   'TN': 40594,\n",
       "   'FN': 636,\n",
       "   'P': 0.30363705343963154,\n",
       "   'R': 0.9999327160987524},\n",
       "  {'TP': 9452428,\n",
       "   'FP': 21713591,\n",
       "   'TN': 3925,\n",
       "   'FN': 56,\n",
       "   'P': 0.3032927625437179,\n",
       "   'R': 0.9999940756313367}]}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c69968f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.6086, device='cuda:0')\n",
      "tensor(0.5814, device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "print(mapi2t)\n",
    "print(mapt2i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "63fddbd1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'pr_i2t': [{'TP': 701,\n",
       "   'FP': 19,\n",
       "   'TN': 21717497,\n",
       "   'FN': 9451783,\n",
       "   'P': 0.9736111111111111,\n",
       "   'R': 7.416040058888225e-05},\n",
       "  {'TP': 7177,\n",
       "   'FP': 463,\n",
       "   'TN': 21717053,\n",
       "   'FN': 9445307,\n",
       "   'P': 0.9393979057591623,\n",
       "   'R': 0.000759271319581181},\n",
       "  {'TP': 66733,\n",
       "   'FP': 7610,\n",
       "   'TN': 21709906,\n",
       "   'FN': 9385751,\n",
       "   'P': 0.8976366302140081,\n",
       "   'R': 0.00705983739300696},\n",
       "  {'TP': 291338,\n",
       "   'FP': 46934,\n",
       "   'TN': 21670582,\n",
       "   'FN': 9161146,\n",
       "   'P': 0.8612536656891495,\n",
       "   'R': 0.030821316386253603},\n",
       "  {'TP': 894223,\n",
       "   'FP': 180538,\n",
       "   'TN': 21536978,\n",
       "   'FN': 8558261,\n",
       "   'P': 0.832020328240418,\n",
       "   'R': 0.09460190570013131},\n",
       "  {'TP': 1934983,\n",
       "   'FP': 470687,\n",
       "   'TN': 21246829,\n",
       "   'FN': 7517501,\n",
       "   'P': 0.8043426571391754,\n",
       "   'R': 0.20470629730767065},\n",
       "  {'TP': 3367698,\n",
       "   'FP': 1015851,\n",
       "   'TN': 20701665,\n",
       "   'FN': 6084786,\n",
       "   'P': 0.768258322195098,\n",
       "   'R': 0.3562765089049609},\n",
       "  {'TP': 4857378,\n",
       "   'FP': 2064380,\n",
       "   'TN': 19653136,\n",
       "   'FN': 4595106,\n",
       "   'P': 0.7017549587835923,\n",
       "   'R': 0.5138731787327013},\n",
       "  {'TP': 6363615,\n",
       "   'FP': 4261808,\n",
       "   'TN': 17455708,\n",
       "   'FN': 3088869,\n",
       "   'P': 0.5989046271381384,\n",
       "   'R': 0.6732214516311268},\n",
       "  {'TP': 7773253,\n",
       "   'FP': 8442979,\n",
       "   'TN': 13274537,\n",
       "   'FN': 1679231,\n",
       "   'P': 0.4793501350992018,\n",
       "   'R': 0.8223502943776472},\n",
       "  {'TP': 8777460,\n",
       "   'FP': 13923633,\n",
       "   'TN': 7793883,\n",
       "   'FN': 675024,\n",
       "   'P': 0.3866536294089452,\n",
       "   'R': 0.9285876601325112},\n",
       "  {'TP': 9268024,\n",
       "   'FP': 18600532,\n",
       "   'TN': 3116984,\n",
       "   'FN': 184460,\n",
       "   'P': 0.3325620459129637,\n",
       "   'R': 0.9804855527922608},\n",
       "  {'TP': 9421005,\n",
       "   'FP': 20985668,\n",
       "   'TN': 731848,\n",
       "   'FN': 31479,\n",
       "   'P': 0.3098334697781635,\n",
       "   'R': 0.9966697642651392},\n",
       "  {'TP': 9449696,\n",
       "   'FP': 21640287,\n",
       "   'TN': 77229,\n",
       "   'FN': 2788,\n",
       "   'P': 0.30394664416509976,\n",
       "   'R': 0.9997050510744054},\n",
       "  {'TP': 9452387,\n",
       "   'FP': 21713968,\n",
       "   'TN': 3548,\n",
       "   'FN': 97,\n",
       "   'P': 0.3032881772667994,\n",
       "   'R': 0.9999897381471368},\n",
       "  {'TP': 9452483,\n",
       "   'FP': 21717497,\n",
       "   'TN': 19,\n",
       "   'FN': 1,\n",
       "   'P': 0.30325598540647125,\n",
       "   'R': 0.9999998942077024}],\n",
       " 'pr_t2i': [{'TP': 501,\n",
       "   'FP': 36,\n",
       "   'TN': 21717480,\n",
       "   'FN': 9451983,\n",
       "   'P': 0.9329608938547486,\n",
       "   'R': 5.300194107707561e-05},\n",
       "  {'TP': 6004,\n",
       "   'FP': 526,\n",
       "   'TN': 21716990,\n",
       "   'FN': 9446480,\n",
       "   'P': 0.919448698315467,\n",
       "   'R': 0.0006351769545444351},\n",
       "  {'TP': 54787,\n",
       "   'FP': 7471,\n",
       "   'TN': 21710045,\n",
       "   'FN': 9397697,\n",
       "   'P': 0.8799993575122875,\n",
       "   'R': 0.00579604260636675},\n",
       "  {'TP': 248570,\n",
       "   'FP': 50790,\n",
       "   'TN': 21666726,\n",
       "   'FN': 9203914,\n",
       "   'P': 0.8303380545163015,\n",
       "   'R': 0.026296791404248874},\n",
       "  {'TP': 759013,\n",
       "   'FP': 210533,\n",
       "   'TN': 21506983,\n",
       "   'FN': 8693471,\n",
       "   'P': 0.7828540368378601,\n",
       "   'R': 0.08029772914717444},\n",
       "  {'TP': 1657182,\n",
       "   'FP': 572005,\n",
       "   'TN': 21145511,\n",
       "   'FN': 7795302,\n",
       "   'P': 0.7434019667259858,\n",
       "   'R': 0.1753170912534737},\n",
       "  {'TP': 2926782,\n",
       "   'FP': 1217829,\n",
       "   'TN': 20499687,\n",
       "   'FN': 6525702,\n",
       "   'P': 0.7061656691062201,\n",
       "   'R': 0.3096309922344222},\n",
       "  {'TP': 4357199,\n",
       "   'FP': 2383086,\n",
       "   'TN': 19334430,\n",
       "   'FN': 5095285,\n",
       "   'P': 0.646441359675444,\n",
       "   'R': 0.4609580931319217},\n",
       "  {'TP': 5912653,\n",
       "   'FP': 4642339,\n",
       "   'TN': 17075177,\n",
       "   'FN': 3539831,\n",
       "   'P': 0.5601759811850165,\n",
       "   'R': 0.6255131455393101},\n",
       "  {'TP': 7446760,\n",
       "   'FP': 8701424,\n",
       "   'TN': 13016092,\n",
       "   'FN': 2005724,\n",
       "   'P': 0.4611515449662947,\n",
       "   'R': 0.7878098497707058},\n",
       "  {'TP': 8613026,\n",
       "   'FP': 13991641,\n",
       "   'TN': 7725875,\n",
       "   'FN': 839458,\n",
       "   'P': 0.3810286610282735,\n",
       "   'R': 0.9111918094756891},\n",
       "  {'TP': 9213223,\n",
       "   'FP': 18557447,\n",
       "   'TN': 3160069,\n",
       "   'FN': 239261,\n",
       "   'P': 0.3317609189839496,\n",
       "   'R': 0.9746880290937282},\n",
       "  {'TP': 9408594,\n",
       "   'FP': 20943085,\n",
       "   'TN': 774431,\n",
       "   'FN': 43890,\n",
       "   'P': 0.30998594838855537,\n",
       "   'R': 0.9953567760601341},\n",
       "  {'TP': 9448378,\n",
       "   'FP': 21626869,\n",
       "   'TN': 90647,\n",
       "   'FN': 4106,\n",
       "   'P': 0.3040483636381072,\n",
       "   'R': 0.9995656168262226},\n",
       "  {'TP': 9452388,\n",
       "   'FP': 21712511,\n",
       "   'TN': 5005,\n",
       "   'FN': 96,\n",
       "   'P': 0.3033023787434703,\n",
       "   'R': 0.9999898439394344},\n",
       "  {'TP': 9452484,\n",
       "   'FP': 21717467,\n",
       "   'TN': 49,\n",
       "   'FN': 0,\n",
       "   'P': 0.30325629963293815,\n",
       "   'R': 1.0}]}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa307ccc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:wuhongfa] *",
   "language": "python",
   "name": "conda-env-wuhongfa-py"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}