Add required folders and files

configs/phase2/FineCapEval_clipRN50_clips_grammar.yml

@@ -0,0 +1,64 @@
+caption_model: transformer
+noamopt: true
+noamopt_warmup: 20000
+label_smoothing: 0.0
+input_json: data/FineCapEval.json
+input_label_h5: none
+input_fc_dir: data/FineCapEval_clip_RN50_fc
+input_att_dir: data/FineCapEval_clip_RN50_att
+input_clipscore_vis_dir: data/FineCapEval_clipscore_vis
+seq_per_img: 5
+batch_size: 160
+learning_rate: 0.0005
+
+checkpoint_path: ./save/clipRN50_clips_grammar/clipRN50_clips_grammar
+
+use_multi_rewards: true
+use_grammar: true
+use_grammar_baseline: true
+# clip_load_path: '/scratch-space/retrieval/save/clip_negative_text/clip_negative_text-epoch=10.ckpt'
+
+# Notice: because I'm to lazy, I reuse the option name for RNNs to set the hyperparameters for transformer:
+# N=num_layers
+# d_model=input_encoding_size
+# d_ff=rnn_size
+
+# will be ignored
+num_layers: 6
+input_encoding_size: 512
+rnn_size: 2048
+
+# Transformer config
+N_enc: 6
+N_dec: 6
+d_model: 512
+d_ff: 2048
+num_att_heads: 8
+dropout: 0.1
+
+
+learning_rate_decay_start: 0
+scheduled_sampling_start: -1 
+save_checkpoint_every: 3000
+language_eval: 0
+val_images_use: 5000
+max_epochs: 15
+train_sample_n: 5
+
+REFORWARD: false
+
+# _BASE_: transformer.yml
+reduce_on_plateau: false
+noamopt: false
+learning_rate: 0.000005
+learning_rate_decay_start: -1
+
+self_critical_after: 15
+max_epochs: 50
+
+verbose: false
+precision: 32
+
+# use_clipscore: true
+use_clipscore: false
+clipscore_reward_weight: 2.0

configs/phase2/clipRN50_clips_grammar.yml

@@ -0,0 +1,64 @@
+caption_model: transformer
+noamopt: true
+noamopt_warmup: 20000
+label_smoothing: 0.0
+input_json: data/cocotalk.json
+input_label_h5: data/cocotalk_label.h5
+input_fc_dir: data/cocotalk_clip_RN50_fc
+input_att_dir: data/cocotalk_clip_RN50_att
+input_clipscore_vis_dir: data/cocotalk_clipscore_vis
+seq_per_img: 5
+batch_size: 160
+learning_rate: 0.0005
+
+checkpoint_path: save/clipRN50_clips_grammar/clipRN50_clips_grammar
+
+use_multi_rewards: true
+use_grammar: true
+use_grammar_baseline: true
+# clip_load_path: '/scratch-space/retrieval/save/clip_negative_text/clip_negative_text-epoch=10.ckpt'
+clip_load_path: 'retrieval/save/clip_negative_text/clip_negative_text-epoch=12.ckpt'
+
+# Notice: because I'm to lazy, I reuse the option name for RNNs to set the hyperparameters for transformer:
+# N=num_layers
+# d_model=input_encoding_size
+# d_ff=rnn_size
+
+# will be ignored
+num_layers: 6
+input_encoding_size: 512
+rnn_size: 2048
+
+# Transformer config
+N_enc: 6
+N_dec: 6
+d_model: 512
+d_ff: 2048
+num_att_heads: 8
+dropout: 0.1
+
+
+learning_rate_decay_start: 0
+scheduled_sampling_start: -1 
+save_checkpoint_every: 3000
+language_eval: 1
+val_images_use: 5000
+max_epochs: 15
+train_sample_n: 5
+
+REFORWARD: false
+
+# _BASE_: transformer.yml
+reduce_on_plateau: false
+noamopt: false
+learning_rate: 0.000005
+learning_rate_decay_start: -1
+
+self_critical_after: 15
+max_epochs: 40
+
+verbose: false
+precision: 32
+
+use_clipscore: true
+clipscore_reward_weight: 2.0

configs/phase2/transformer.yml

@@ -0,0 +1,41 @@
+caption_model: transformer
+noamopt: true
+noamopt_warmup: 20000
+label_smoothing: 0.0
+input_json: data/cocotalk.json
+input_label_h5: data/cocotalk_label.h5
+input_att_dir: data/cocotalk_att
+seq_per_img: 5
+batch_size: 10
+learning_rate: 0.0005
+
+checkpoint_path: ./save/trans_rn50_sc
+
+# Notice: because I'm to lazy, I reuse the option name for RNNs to set the hyperparameters for transformer:
+# N=num_layers
+# d_model=input_encoding_size
+# d_ff=rnn_size
+
+# will be ignored
+num_layers: 6
+input_encoding_size: 512
+rnn_size: 2048
+
+# Transformer config
+N_enc: 6
+N_dec: 6
+d_model: 512
+d_ff: 2048
+num_att_heads: 8
+dropout: 0.1
+
+
+learning_rate_decay_start: 0
+scheduled_sampling_start: -1 
+save_checkpoint_every: 3000
+language_eval: 1
+val_images_use: 5000
+max_epochs: 15
+train_sample_n: 5
+
+REFORWARD: false

data/README.md

@@ -0,0 +1 @@
+directory to store preprocessed files

retrieval/README.md

@@ -0,0 +1,5 @@
+# Finetuning CLIP reward model
+
+```bash
+python train_pl.py --cfg clip_negative_text --id clip_negative_text
+```

retrieval/caption_data.py

@@ -0,0 +1,500 @@
+from torch.utils.data import DataLoader, Dataset, Sampler
+from pathlib import Path
+import json
+from multiprocessing import Pool
+from tqdm import tqdm
+from PIL import Image
+import random
+import numpy as np
+import torch
+import torchvision
+import torchvision.transforms as T
+
+from torch.utils.data.distributed import DistributedSampler
+
+from transformers import T5Tokenizer, BertTokenizer, BertTokenizerFast, CLIPTokenizer
+
+import text_utils
+
+project_dir = Path(__file__).parent.resolve()
+workspace_dir = project_dir.parent.parent
+dataset_dir = workspace_dir.joinpath('datasets/').resolve()
+# coco_dir = dataset_dir.joinpath('COCO')
+# vg_dir = dataset_dir.joinpath('VG')
+coco_img_dir = dataset_dir.joinpath('COCO/images/')
+coco_data_dir = project_dir.parent.joinpath('CLIP-ViL/CLIP-ViL-Direct/caption/data/')
+# coco_feature_dir = coco_dir.joinpath('features')
+
+
+class COCORetrievalDataset(Dataset):
+    def __init__(self, split='karpathy_train', rank=-1, topk=-1, verbose=True, args=None, mode='train'):
+        super().__init__()
+
+        self.topk = topk
+        self.verbose = verbose
+        self.args = args
+        self.rank = rank
+        self.mode = mode
+
+        # Loading datasets to data
+        self.source = split
+        if self.verbose:
+            print('Data source: ', self.source)
+
+        # if self.args.tokenizer is None:
+        #     self.args.tokenizer = self.args.decoder_backbone
+
+        # if 'bert' in self.args.tokenizer:
+        #     self.tokenizer = BertTokenizerFast.from_pretrained(
+        #         self.args.tokenizer,
+        #         # max_length=self.args.max_text_length,
+        #         # do_lower_case=self.args.do_lower_case
+        #         )
+        # elif 'clip' in self.args.tokenizer:
+        #     self.tokenizer = CLIPTokenizer.from_pretrained(
+        #         self.args.tokenizer,
+        #         # max_length=self.args.max_text_length,
+        #         # do_lower_case=self.args.do_lower_case
+        #         )
+
+        self.tokenizer = CLIPTokenizer.from_pretrained(
+                self.args.tokenizer,
+                # max_length=self.args.max_text_length,
+                # do_lower_case=self.args.do_lower_case
+                )
+
+        with open(coco_data_dir.joinpath('cocotalk.json')) as f:
+            self.vocab = list(json.load(f)['ix_to_word'].values())
+            popped = self.vocab.pop(-1)
+            assert popped == 'UNK'
+            if self.verbose:
+                print('vocab size: ', len(self.vocab))
+
+
+        data_info_path = coco_data_dir.joinpath('dataset_coco.json')
+        with open(data_info_path) as f:
+            karpathy_data = json.load(f)
+
+        split_rename = {
+            'train': 'train',
+            'restval': 'train',
+            'val': 'val',
+            'test': 'test'
+        }
+
+        n_images = 0
+
+        data = []
+        # self.vocab = set()
+        for datum in karpathy_data['images']:
+            re_split = split_rename[datum['split']]
+
+            # if re_split == 'train':
+            #     for d in datum['sentences']:
+            #         self.vocab = self.vocab.union(set(d['tokens']))
+
+            if re_split != self.source.split('_')[-1]:
+                continue
+
+            if re_split == 'train':
+                # for d in datum['sentences']:
+                #     img_id = datum['filename'].split('.')[0]
+                #     new_datum = {
+                #         'filename': datum['filename'],
+                #         'img_id': img_id,
+                #         'sent': d['raw'].strip(),
+                #         'targets': [d['raw'].strip() for d in datum['sentences']],
+                #         'is_train': True,
+                #         'cocoid': datum['cocoid']
+                #     }
+                #     data.append(new_datum)
+                img_id = datum['filename'].split('.')[0]
+                new_datum = {
+                    'filename': datum['filename'],
+                    'img_id': img_id,
+                    # 'sent': d['raw'],
+                    # 'targets': [d['raw'].strip() for d in datum['sentences']],
+                    'targets': [" ".join(d['tokens']) for d in datum['sentences']],
+                    'is_train': True,
+                    'cocoid': datum['cocoid']
+                }
+                data.append(new_datum)
+
+            else:
+                img_id = datum['filename'].split('.')[0]
+                new_datum = {
+                    'filename': datum['filename'],
+                    'img_id': img_id,
+                    # 'sent': d['raw'],
+                    # 'targets': [d['raw'].strip() for d in datum['sentences']],
+                    'targets': [" ".join(d['tokens']) for d in datum['sentences']],
+                    'is_train': False,
+                    'cocoid': datum['cocoid']
+                }
+                data.append(new_datum)
+
+            n_images += 1
+
+        if self.verbose:
+            print(f"{self.source} has {n_images} images")
+            # print(f"Loaded {len(data)} data from", split)
+
+        self.n_gpus = torch.cuda.device_count()
+
+        if self.topk > 0:
+            data = data[:self.topk]
+            if self.verbose:
+                print(f"Use only {self.topk} data")
+
+        self.data = data
+
+        # if self.verbose:
+        #     print("# all sentences:", len(self.data))
+
+        if self.args.load_feat:
+            # feat_dir = coco_dir.joinpath(''
+            # self.feat_loader = HybridLoader('/scratch-space/CLIP-ViL/CLIP-ViL-Direct/caption/data/cocotalk_clipscore_vis', ext='.npy', in_memory=False)
+            self.feat_loader = HybridLoader(
+                coco_data_dir.joinpath('cocotalk_clipscore_vis'),
+                ext='.npy', in_memory=False)
+        else:
+            if 'openai/clip' in self.args.encoder_backbone:
+                # from transformers import CLIPProcessor
+                # self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32",
+                #     size=args.image_size,
+                #     do_resize=True,
+                #     do_center_crop=False,
+                # )
+                # self.img_transform = lambda image: self.processor.feature_extractor(
+                #     image,
+                #     return_tensors='pt')['pixel_values'][0]
+
+                self.image_mean = [0.48145466, 0.4578275, 0.40821073]
+                self.image_std = [0.26862954, 0.26130258, 0.27577711]
+
+                # captioning
+                # self.img_transform = T.Compose([
+                #     T.Resize((self.args.image_size, self.args.image_size))
+                # ])
+
+                # retrieval
+                self.img_transform = T.Compose([
+                    T.Resize(self.args.image_size, interpolation=T.functional.InterpolationMode.BICUBIC),
+                    T.CenterCrop(self.args.image_size)
+                ])
+
+                self.img_tensor_transform = T.Compose([
+                    # T.RandomCrop(224),
+                    # T.RandomHorizontalFlip(p=0.3),
+                    T.ConvertImageDtype(torch.float),
+                    T.Normalize(self.image_mean, self.image_std)
+                ]
+                )
+            # elif 'google/vit' in self.args.encoder_backbone:
+            #     self.image_mean = [0.5, 0.5, 0.5]
+            #     self.image_std = [0.5, 0.5, 0.5]
+
+            #     self.img_transform = T.Compose([
+            #         # T.PILToTensor(),
+            #         T.Resize((self.args.image_size, self.args.image_size))
+            #     ])
+
+            #     self.img_tensor_transform = T.Compose([
+            #         # T.RandomCrop(224),
+            #         # T.RandomHorizontalFlip(p=0.3),
+            #         T.ConvertImageDtype(torch.float),
+            #         T.Normalize(self.image_mean, self.image_std)
+            #     ]
+            #     )
+
+    def get_negative_text(self, text):
+        neg_type = random.choice(['repeat', 'remove', 'insert', 'swap', 'shuffle'])
+
+        if neg_type == 'repeat':
+            text = text_utils.repeat(text)
+        elif neg_type == 'remove':
+            text = text_utils.remove(text)
+        elif neg_type == 'insert':
+            text = text_utils.insert(text, self.vocab)
+        elif neg_type == 'swap':
+            text = text_utils.swap(text, self.vocab)
+        elif neg_type == 'shuffle':
+            text = text_utils.shuffle(text)
+
+        return text, neg_type
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        datum = self.data[idx]
+        return self.process_datum(datum)
+
+    def process_datum(self, datum):
+        out_dict = {}
+
+        ###### Image ######
+
+        if self.args.load_feat:
+            cocoid = datum['cocoid']
+            out_dict['cocoid'] = str(cocoid)
+            img_feat = self.feat_loader.get(str(cocoid))
+            out_dict['img_feat'] = torch.from_numpy(img_feat)
+
+        else:
+            img_id = datum['img_id']
+            out_dict['img_id'] = img_id
+
+            if 'train' in datum['filename']:
+                img_split = 'train2014'
+            elif 'val' in datum['filename']:
+                img_split = 'val2014'
+            img_path = coco_img_dir.joinpath(img_split).joinpath(datum['filename']).with_suffix('.jpg')
+            assert img_path.exists()
+            img_path = str(img_path)
+            out_dict['img_path'] = img_path
+
+            img_tensor = torchvision.io.read_image(img_path)
+            # out_dict['img_tensor'] = img
+
+            # img = Image.open(img_path).convert('RGB')
+            # img_tensor = torch.as_tensor(np.asarray(img))
+            out_dict['img_tensor'] = self.img_transform(img_tensor)
+            # self.img_transform(img_tensor)
+            # out_dict['img_tensor'] = self.img_transform(img)
+
+        ###### Text #####
+        # if datum['is_train']:
+        # sent = datum['sent'].strip()
+
+        sent = random.choice(datum['targets'])
+        
+        # target_ids = self.tokenizer.encode(
+        #     sent, max_length=self.args.gen_max_length, truncation=True)
+
+        # assert len(target_ids) <= self.args.gen_max_length, len(target_ids)
+        out_dict['sent'] = sent
+        # out_dict['target_ids'] = torch.LongTensor(target_ids)
+        # out_dict['target_length'] = len(target_ids)
+
+
+        # negative sample
+        neg_sent, neg_type = self.get_negative_text(sent)
+
+        # neg_target_ids = self.tokenizer.encode(
+        #     neg_sent, max_length=self.args.gen_max_length, truncation=True)
+
+        # assert len(neg_target_ids) <= self.args.gen_max_length, len(neg_target_ids)
+        out_dict['neg_sent'] = neg_sent
+        out_dict['neg_type'] = neg_type
+        # out_dict['neg_target_ids'] = torch.LongTensor(neg_target_ids)
+        # out_dict['neg_target_length'] = len(neg_target_ids)
+
+
+        if 'targets' in datum:
+            out_dict['targets'] = datum['targets']
+
+        return out_dict
+
+    def collate_fn(self, batch):
+        batch_entry = {}
+
+        B = len(batch)
+
+        # if 'target_ids' in batch[0]:
+        #     T_W_L = max(entry['target_length'] for entry in batch)
+        #     target_ids = torch.ones(
+        #         B, T_W_L, dtype=torch.long) * self.tokenizer.pad_token_id
+
+        # if 'target_ids' in batch[0]:
+        #     T_W_L = max(entry['target_length'] for entry in batch)
+        #     target_ids = torch.ones(
+        #         B, T_W_L, dtype=torch.long) * self.tokenizer.pad_token_id
+
+
+
+        targets = []
+        img_ids = []
+        img_paths = []
+
+        coco_ids = []
+
+        if self.args.load_feat:
+            img_feats = torch.zeros(B, 512, dtype=torch.float)
+        else:
+            # imgs = []
+            img_tensor = torch.zeros(B, 3, self.args.image_size, self.args.image_size, dtype=torch.uint8)
+
+        for i, entry in enumerate(batch):
+
+            if self.args.load_feat:
+                coco_ids.append(entry['cocoid'])
+                img_feats[i] = entry['img_feat']
+
+            else:
+
+                img_ids.append(entry['img_id'])
+                img_paths.append(entry['img_path'])
+                img_tensor[i] = entry['img_tensor']
+
+            # if 'target_ids' in entry:
+            #     target_ids[i, :entry['target_length']] = entry['target_ids']
+
+            if 'targets' in entry:
+                targets.append(entry['targets'])
+
+        if 'sent' in batch[0]:
+            # word_mask = target_ids != self.tokenizer.pad_token_id
+            # target_ids[~word_mask] = -100
+            # batch_entry['target_ids'] = target_ids
+
+            tokenized = self.tokenizer([entry['sent'] for entry in batch], truncation=True, padding=True, return_tensors='pt')
+            neg_tokenized = self.tokenizer([entry['neg_sent'] for entry in batch], truncation=True, padding=True, return_tensors='pt')
+                #     sent, max_length=self.args.gen_max_length, truncation=True)
+
+            batch_entry['text'] = (tokenized.input_ids, tokenized.attention_mask)
+            batch_entry['neg_text'] = (neg_tokenized.input_ids, neg_tokenized.attention_mask)
+
+
+        if self.args.load_feat:
+            batch_entry['coco_ids'] = coco_ids
+            batch_entry['img_feats'] = img_feats
+
+        else:
+
+            img_tensor = self.img_tensor_transform(img_tensor)
+
+            batch_entry['img_id'] = img_ids
+            batch_entry['img_paths'] = img_paths
+            batch_entry['img_tensor'] = img_tensor
+
+        batch_entry['targets'] = targets
+
+        # print('batch created')
+
+        # batch_entry['task'] = 'caption'
+
+        return batch_entry
+
+
+# def get_loader(args, split='karpathy_train', mode='train',
+#                batch_size=32, workers=4, distributed=False, gpu=0,
+#                topk=-1):
+
+#     verbose = (gpu == 0)
+
+#     dataset = COCORetrievalDataset(
+#         split,
+#         rank=gpu,
+#         topk=topk,
+#         verbose=verbose,
+#         args=args,
+#         mode=mode)
+
+#     # if distributed:
+#     #     sampler = DistributedSampler(dataset)
+#     # else:
+#     #     sampler = None
+
+#     if mode == 'train':
+#         loader = DataLoader(
+#             dataset, batch_size=batch_size, shuffle=(sampler is None),
+#             num_workers=workers, pin_memory=True, sampler=sampler,
+#             collate_fn=dataset.collate_fn)
+#     else:
+#         loader = DataLoader(
+#             dataset,
+#             batch_size=batch_size, shuffle=False,
+#             num_workers=workers, pin_memory=True,
+#             sampler=sampler,
+#             collate_fn=dataset.collate_fn,
+#             drop_last=False)
+
+#     # if verbose:
+#         # loader.evaluator = COCOCaptionEvaluator()
+
+#     # loader.task = 'caption'
+
+#     return loader
+
+
+# class COCOCaptionEvaluator:
+#     def __init__(self):
+#         import language_evaluation
+#         self.evaluator = language_evaluation.CocoEvaluator(verbose=False)
+
+#     def evaluate(self, predicts, answers):
+
+#         results = self.evaluator.run_evaluation(predicts, answers)
+
+#         return results
+
+import six
+import os
+import h5py
+
+class HybridLoader:
+    """
+    If db_path is a director, then use normal file loading
+    If lmdb, then load from lmdb
+    The loading method depend on extention.
+
+    in_memory: if in_memory is True, we save all the features in memory
+               For individual np(y|z)s, we don't need to do that because the system will do this for us.
+               Should be useful for lmdb or h5.
+               (Copied this idea from vilbert)
+    """
+
+    def __init__(self, db_path, ext='.npy', in_memory=False):
+        self.db_path = db_path
+        self.ext = ext
+        if self.ext == '.npy':
+            self.loader = lambda x: np.load(six.BytesIO(x))
+        else:
+            self.loader = lambda x: np.load(six.BytesIO(x))['feat']
+        # if db_path.endswith('.lmdb'):
+        #     self.db_type = 'lmdb'
+        #     self.lmdb = lmdbdict(db_path, unsafe=True)
+        #     self.lmdb._key_dumps = DUMPS_FUNC['ascii']
+        #     self.lmdb._value_loads = LOADS_FUNC['identity']
+        # elif db_path.endswith('.pth'):  # Assume a key,value dictionary
+        #     self.db_type = 'pth'
+        #     self.feat_file = torch.load(db_path)
+        #     self.loader = lambda x: x
+        #     print('HybridLoader: ext is ignored')
+        # elif db_path.endswith('h5'):
+        #     self.db_type = 'h5'
+        #     self.loader = lambda x: np.array(x).astype('float32')
+        # else:
+        #     self.db_type = 'dir'
+
+        self.in_memory = in_memory
+        if self.in_memory:
+            self.features = {}
+
+    def get(self, key):
+
+        # if self.in_memory and key in self.features:
+        #     # We save f_input because we want to save the
+        #     # compressed bytes to save memory
+        #     f_input = self.features[key]
+        # elif self.db_type == 'lmdb':
+        #     f_input = self.lmdb[key]
+        # elif self.db_type == 'pth':
+        #     f_input = self.feat_file[key]
+        # elif self.db_type == 'h5':
+        #     f_input = h5py.File(self.db_path, 'r')[key]
+        # else:
+            # f_input = open(os.path.join(
+            #     self.db_path, key + self.ext), 'rb').read()
+
+        f_input = open(os.path.join(
+            self.db_path, key + self.ext), 'rb').read()
+
+        if self.in_memory and key not in self.features:
+            self.features[key] = f_input
+
+        # load image
+        feat = self.loader(f_input)
+
+        return feat

retrieval/clip_model.py

@@ -0,0 +1,350 @@
+from transformers import CLIPModel, CLIPTokenizer
+import os
+import json
+import argparse
+from random import shuffle, seed
+import string
+# non-standard dependencies:
+import h5py
+from six.moves import cPickle
+import numpy as np
+import torch
+import torchvision.models as models
+import skimage.io
+
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+from PIL import Image
+from torch import nn
+
+
+class CLIPScore(nn.Module):
+    def __init__(self, clipscore_w=2.5, image_size=224, mode='clip_s', use_grammar=False, joint_out=False):
+        super(CLIPScore, self).__init__()
+        # from transformers import CLIPModel, CLIPTokenizer
+        self.clip_model = CLIPModel.from_pretrained(
+            'openai/clip-vit-base-patch32')
+        self.tokenizer = CLIPTokenizer.from_pretrained(
+            'openai/clip-vit-base-patch32')
+
+        self.clip_model.eval()
+
+        self.clipscore_w = clipscore_w
+
+        self.image_transform = self._transform(image_size)
+
+        self.mode = mode
+        assert mode in ['clip_s', 'refclip_s']
+
+        self.use_grammar = use_grammar
+        self.joint_out = joint_out
+
+        if self.use_grammar and self.joint_out is False:
+            self.grammar_score_head = nn.Sequential(
+                nn.Linear(self.clip_model.text_embed_dim, self.clip_model.projection_dim, bias=False),
+                nn.ReLU(),
+                nn.Linear(self.clip_model.projection_dim, 2, bias=False)
+            )
+
+    def _transform(self, n_px):
+        return Compose([
+            Resize(n_px, interpolation=Image.BICUBIC),
+            CenterCrop(n_px),
+            lambda image: image.convert("RGB"),
+            ToTensor(),
+            Normalize((0.48145466, 0.4578275, 0.40821073),
+                      (0.26862954, 0.26130258, 0.27577711)),
+        ])
+
+    def load_image(self, image_path):
+        image = Image.open(image_path)
+        return image
+
+    # @torch.no_grad()
+    def image_extract(self, image):
+        if isinstance(image, str):
+            image = self.load_image(image)
+        if not isinstance(image, torch.Tensor):
+            image = self.image_transform(image)
+
+        img_tensor = image.view(-1, 3, 224, 224)
+        device = next(self.clip_model.parameters()).device
+        img_tensor = img_tensor.to(device)
+
+        clip_model = self.clip_model
+
+        img_feat = clip_model.vision_model(img_tensor).pooler_output
+        img_feat = clip_model.visual_projection(img_feat)
+        img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)
+
+        return img_feat
+
+    # @torch.no_grad()
+    def text_extract(self, text, prompt="A photo depicts", proj_norm=True):
+        if isinstance(text, str):
+            text_batch = [" ".join([prompt, text])]
+        elif isinstance(text, list):
+            text_batch = [" ".join([prompt, txt]) for txt in text]
+        
+        if isinstance(text, tuple) and isinstance(text[0], torch.Tensor):
+            input_ids, attention_mask = text
+        else:
+            input_text = text_batch
+
+            tokenized = self.tokenizer(
+                input_text, return_tensors='pt', padding=True)
+
+            input_ids = tokenized.input_ids
+            attention_mask = tokenized.attention_mask
+
+        clip_model = self.clip_model
+        device = next(self.clip_model.parameters()).device
+        input_ids = input_ids.to(device)
+        attention_mask = attention_mask.to(device)
+
+        text_feat = clip_model.text_model(input_ids, attention_mask).pooler_output
+
+        if proj_norm:
+            text_feat = clip_model.text_projection(text_feat)
+            text_feat = text_feat / text_feat.norm(dim=-1, keepdim=True)
+
+        return text_feat
+
+    # @torch.no_grad()
+    def calc_clip_s(self, img_feat, text_feat):
+        return self.clipscore_w * torch.relu((img_feat * text_feat).sum(dim=-1))
+
+    # @torch.no_grad()
+    def calc_refclip_s(self, img_feat=None, text_feat=None, ref_text_feat=None, ref_text_mask=None, clip_s=None):
+
+        if clip_s is None:
+            clip_s = self.calc_clip_s(img_feat, text_feat)
+
+        B, dim = img_feat.size()
+
+        ref_text_feat = ref_text_feat.view(B, -1, dim)
+
+        K = ref_text_feat.size(1)
+
+        text_feat = text_feat.view(B, 1, dim).expand(-1, K, -1)
+        assert ref_text_feat.size() == text_feat.size(
+        ), (ref_text_feat.size(), text_feat.size())
+
+        ref_score = self.calc_clip_s(text_feat, ref_text_feat)
+        if ref_text_mask is not None:
+            if not isinstance(ref_text_mask, torch.Tensor):
+                ref_text_mask = torch.tensor(
+                    ref_text_mask, dtype=ref_score.dtype, device=ref_score.device)
+            ref_score = ref_score.view(B, K) * ref_text_mask.view(B, K)
+
+        ref_score = ref_score.view(B, K).max(dim=1).values
+
+        assert clip_s.size() == (B,)
+        assert clip_s.size() == ref_score.size()
+
+        # harmonic mean
+        refclip_s = 2 / (1 / clip_s + 1 / ref_score)
+        return refclip_s
+
+    # # @torch.no_grad()
+    # def forward(self,
+    #             images=None, text=None,
+    #             img_feat=None, text_feat=None,
+    #             ref_text=None, ref_text_feat=None, ref_text_mask=None,
+    #             prompt="A photo depicts",
+    #             mode=None):
+    #     if img_feat is None:
+    #         img_feat = self.image_extract(images)
+    #     img_feat = img_feat.view(-1, 512)
+
+    #     if text_feat is None:
+    #         text_feat = self.text_extract(text, prompt=prompt)
+    #     text_feat = text_feat.view(-1, 512)
+
+    #     if mode is None:
+    #         mode = self.mode
+    #     assert mode in ['clip_s', 'refclip_s']
+
+    #     if mode == 'clip_s':
+    #         clip_s = self.calc_clip_s(img_feat, text_feat)
+    #         return clip_s
+    #     elif mode == 'refclip_s':
+    #         if ref_text_feat is None:
+    #             ref_text_feat = self.text_extract(ref_text, prompt=prompt)
+    #         ref_text_feat = ref_text_feat.view(-1, 512)
+
+    #         refclip_s = self.calc_refclip_s(
+    #             img_feat, text_feat, ref_text_feat, ref_text_mask=ref_text_mask)
+    #         return refclip_s
+
+
+    def train_step(self,
+                   images=None, text=None,
+                   img_feat=None, text_feat=None,
+                   neg_text=None, neg_text_feat=None,
+                #    ref_text=None, ref_text_feat=None, ref_text_mask=None,
+                   prompt="A photo depicts",
+                #    return_loss=True,
+                   **kwargs):
+
+        if img_feat is None:
+            img_feat = self.image_extract(images)
+        img_feat = img_feat.view(-1, 512)
+
+        B = img_feat.size(0)
+
+        if self.joint_out:
+            pos_text_feat = self.text_extract(text, prompt=prompt, proj_norm=False).view(B, 512)
+            neg_text_feat = self.text_extract(neg_text, prompt=prompt, proj_norm=False).view(-1, 512)
+            neg_B = neg_text_feat.size(0)
+
+            # [B+neg_B, 512]
+            text_feat = torch.cat([pos_text_feat, neg_text_feat], dim=0)
+
+            text_cont_feat = self.clip_model.text_projection(text_feat)
+            text_cont_feat = text_cont_feat / text_cont_feat.norm(dim=-1, keepdim=True)
+
+            text_cont_feat = text_cont_feat.view(B+neg_B, 512)
+
+            logit_scale = self.clip_model.logit_scale.exp()
+
+            # [B+neg_B * B]
+            logits_per_text = torch.matmul(text_cont_feat, img_feat.t()) * logit_scale
+
+            # image-to-text label: positive text
+            caption_loss = -torch.diag(nn.functional.log_softmax(logits_per_text, dim=0)[:B]).mean()
+
+            # calculate text-to-image only on positive text 
+            image_loss = -torch.diag(nn.functional.log_softmax(logits_per_text[:B], dim=1)).mean()
+
+            clip_loss = (caption_loss + image_loss) / 2.0
+
+            out = {
+                'clip_loss': clip_loss,
+                'img_feat': img_feat,
+                'text_feat': text_cont_feat[:B].detach(),
+                # 'neg_text_feat': neg_text_feat,
+            }
+
+            return out
+
+
+        else:
+            if text_feat is None:
+                text_feat = self.text_extract(text, prompt=prompt, proj_norm=False)
+
+            text_cont_feat = self.clip_model.text_projection(text_feat)
+            text_cont_feat = text_cont_feat / \
+                text_cont_feat.norm(dim=-1, keepdim=True)
+
+            text_cont_feat = text_cont_feat.view(B, 512)
+
+
+            # cosine similarity as logits
+            logit_scale = self.clip_model.logit_scale.exp()
+            logits_per_text = torch.matmul(text_cont_feat, img_feat.t()) * logit_scale
+            # logits_per_image = logits_per_text.T
+
+            clip_loss = clip_loss_fn(logits_per_text)
+
+
+            # negative sampling
+            pos_text_feat = text_feat.view(B, 512)
+            neg_text_feat = self.text_extract(neg_text, prompt=prompt, proj_norm=False).view(B, 512)
+
+            grammar_text_feat = torch.cat([pos_text_feat, neg_text_feat], dim=0)
+
+            # 2B, 1
+            grammar_text_logit = self.grammar_score_head(grammar_text_feat)
+            grammar_labels = torch.LongTensor([1] * B + [0] * B).to(grammar_text_logit.device).view(2 * B)
+
+            grammar_loss = torch.nn.functional.cross_entropy(grammar_text_logit, grammar_labels)
+
+            grammar_pred = grammar_text_logit.argmax(dim=1, keepdim=False)
+            grammar_pos_pred = grammar_pred[:B]
+            grammar_neg_pred = grammar_pred[B:]
+            # grammar_acc = (grammar_pred == grammar_labels).float().mean()
+
+            out = {
+                'clip_loss': clip_loss,
+                'grammar_loss': grammar_loss,
+                'img_feat': img_feat,
+                'text_feat': text_cont_feat,
+                'neg_text_feat': neg_text_feat,
+                'grammar_pos_pred': grammar_pos_pred,
+                'grammar_neg_pred': grammar_neg_pred,
+            }
+
+            return out
+
+    def train_step_old(self,
+                   images=None, text=None,
+                   img_feat=None, text_feat=None,
+                   neg_text=None, neg_text_feat=None,
+                #    ref_text=None, ref_text_feat=None, ref_text_mask=None,
+                   prompt="A photo depicts",
+                #    return_loss=True,
+                   **kwargs):
+
+        if img_feat is None:
+            img_feat = self.image_extract(images)
+        img_feat = img_feat.view(-1, 512)
+
+        B = img_feat.size(0)
+
+        
+
+        if text_feat is None:
+            text_feat = self.text_extract(text, prompt=prompt, proj_norm=False)
+
+            text_cont_feat = self.clip_model.text_projection(text_feat)
+            text_cont_feat = text_cont_feat / text_cont_feat.norm(dim=-1, keepdim=True)
+        text_cont_feat = text_cont_feat.view(B, 512)
+
+        # cosine similarity as logits
+        logit_scale = self.clip_model.logit_scale.exp()
+        logits_per_text = torch.matmul(text_cont_feat, img_feat.t()) * logit_scale
+        # logits_per_image = logits_per_text.T
+
+        clip_loss = clip_loss_fn(logits_per_text)
+
+
+        # negative sampling
+        pos_text_feat = text_feat.view(B, 512)
+        neg_text_feat = self.text_extract(neg_text, prompt=prompt, proj_norm=False).view(B, 512)
+
+        grammar_text_feat = torch.cat([pos_text_feat, neg_text_feat], dim=0)
+
+        # 2B, 1
+        grammar_text_logit = self.grammar_score_head(grammar_text_feat)
+        grammar_labels = torch.LongTensor([1] * B + [0] * B).to(grammar_text_logit.device).view(2 * B)
+
+        grammar_loss = torch.nn.functional.cross_entropy(grammar_text_logit, grammar_labels)
+
+        grammar_pred = grammar_text_logit.argmax(dim=1, keepdim=False)
+        grammar_pos_pred = grammar_pred[:B]
+        grammar_neg_pred = grammar_pred[B:]
+        # grammar_acc = (grammar_pred == grammar_labels).float().mean()
+
+        out = {
+            'clip_loss': clip_loss,
+            'grammar_loss': grammar_loss,
+            'img_feat': img_feat,
+            'text_feat': text_cont_feat,
+            'neg_text_feat': neg_text_feat,
+            'grammar_pos_pred': grammar_pos_pred,
+            'grammar_neg_pred': grammar_neg_pred,
+        }
+
+        return out
+
+# contrastive loss function, adapted from
+# https://sachinruk.github.io/blog/pytorch/pytorch%20lightning/loss%20function/gpu/2021/03/07/CLIP.html
+def contrastive_loss(logits: torch.Tensor, dim: int) -> torch.Tensor:
+    neg_ce = torch.diag(nn.functional.log_softmax(logits, dim=dim))
+    return -neg_ce.mean()
+
+
+def clip_loss_fn(similarity: torch.Tensor) -> torch.Tensor:
+    caption_loss = contrastive_loss(similarity, dim=0)
+    image_loss = contrastive_loss(similarity, dim=1)
+    return (caption_loss + image_loss) / 2.0

retrieval/configs/clip_negative_text.yaml

@@ -0,0 +1,14 @@
+checkpoint_dir: ./save/clip_negative_text/
+
+losses_log_every: 25
+precision: 32
+load_feat: true
+data_in_memory: false
+
+batch_size: 1600
+valid_batch_size: 200
+clip_grad_norm: 0
+
+epochs: 30
+use_grammar: true
+joint_out: false

retrieval/param.py

@@ -0,0 +1,209 @@
+import argparse
+import random
+
+import numpy as np
+import torch
+
+import pprint
+import yaml
+
+
+def str2bool(v):
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+
+def is_interactive():
+    import __main__ as main
+    return not hasattr(main, '__file__')
+
+
+def get_optimizer(optim, verbose=False):
+    # Bind the optimizer
+    if optim == 'rms':
+        if verbose:
+            print("Optimizer: Using RMSProp")
+        optimizer = torch.optim.RMSprop
+    elif optim == 'adam':
+        if verbose:
+            print("Optimizer: Using Adam")
+        optimizer = torch.optim.Adam
+    elif optim == 'adamw':
+        if verbose:
+            print("Optimizer: Using AdamW")
+        # optimizer = torch.optim.AdamW
+        optimizer = 'adamw'
+    elif optim == 'adamax':
+        if verbose:
+            print("Optimizer: Using Adamax")
+        optimizer = torch.optim.Adamax
+    elif optim == 'sgd':
+        if verbose:
+            print("Optimizer: SGD")
+        optimizer = torch.optim.SGD
+    else:
+        assert False, "Please add your optimizer %s in the list." % optim
+
+    return optimizer
+
+
+def parse_args(parse=True, **optional_kwargs):
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--seed', type=int, default=9595, help='random seed')
+
+    # Data Splits
+    parser.add_argument("--train", default='karpathy_train')
+    parser.add_argument("--valid", default='karpathy_val')
+    parser.add_argument("--test", default='karpathy_test')
+    # parser.add_argument('--test_only', action='store_true')
+
+    # Quick experiments
+    parser.add_argument('--train_topk', type=int, default=-1)
+    parser.add_argument('--valid_topk', type=int, default=-1)
+
+    # Checkpoint
+    parser.add_argument('--output', type=str, default='snap/test')
+    parser.add_argument('--load', type=str, default=None, help='Load the model (usually the fine-tuned model).')
+    parser.add_argument('--from_scratch', action='store_true')
+
+    # CPU/GPU
+    parser.add_argument("--multiGPU", action='store_const', default=False, const=True)
+    parser.add_argument('--fp16', action='store_true')
+    parser.add_argument("--distributed", action='store_true')
+    parser.add_argument("--num_workers", default=0, type=int)
+    parser.add_argument('--local_rank', type=int, default=-1)
+    # parser.add_argument('--rank', type=int, default=-1)
+
+    # Model Config
+    # parser.add_argument('--encoder_backbone', type=str, default='openai/clip-vit-base-patch32')
+    # parser.add_argument('--decoder_backbone', type=str, default='bert-base-uncased')
+    parser.add_argument('--tokenizer', type=str, default='openai/clip-vit-base-patch32')
+
+    # parser.add_argument('--position_embedding_type', type=str, default='absolute')
+
+    # parser.add_argument('--encoder_transform', action='store_true')
+
+    parser.add_argument('--max_text_length', type=int, default=40)
+
+    # parser.add_argument('--image_size', type=int, default=224)
+    # parser.add_argument('--patch_size', type=int, default=32)
+
+    # parser.add_argument('--decoder_num_layers', type=int, default=12)
+
+    # Training
+    parser.add_argument('--batch_size', type=int, default=256)
+    parser.add_argument('--valid_batch_size', type=int, default=None)
+
+    parser.add_argument('--optim', default='adamw')
+
+    parser.add_argument('--warmup_ratio', type=float, default=0.05)
+    parser.add_argument('--weight_decay', type=float, default=0.01)
+    parser.add_argument('--clip_grad_norm', type=float, default=-1.0)
+    parser.add_argument('--gradient_accumulation_steps', type=int, default=1)
+    parser.add_argument('--lr', type=float, default=1e-4)
+    parser.add_argument('--adam_eps', type=float, default=1e-6)
+    parser.add_argument('--adam_beta1', type=float, default=0.9)
+    parser.add_argument('--adam_beta2', type=float, default=0.999)
+
+    parser.add_argument('--epochs', type=int, default=20)
+    # parser.add_argument('--dropout', type=float, default=0.1)
+
+
+    # Inference
+    # parser.add_argument('--num_beams', type=int, default=1)
+    # parser.add_argument('--gen_max_length', type=int, default=20)
+
+    parser.add_argument('--start_from', type=str, default=None)
+
+    # Data
+    # parser.add_argument('--do_lower_case', type=str2bool, default=None)
+
+    # parser.add_argument('--prefix', type=str, default=None)
+
+
+    # COCO Caption
+    # parser.add_argument('--no_prefix', action='store_true')
+
+    parser.add_argument('--no_cls', action='store_true')
+
+    parser.add_argument('--cfg', type=str, default=None)
+    parser.add_argument('--id', type=str, default=None)
+
+    # Etc.
+    parser.add_argument('--comment', type=str, default='')
+    parser.add_argument("--dry", action='store_true')
+
+    # Parse the arguments.
+    if parse:
+        args = parser.parse_args()
+    # For interative engironmnet (ex. jupyter)
+    else:
+        args = parser.parse_known_args()[0]
+
+    loaded_kwargs = {}
+    if args.cfg is not None:
+        cfg_path = f'configs/{args.cfg}.yaml'
+        with open(cfg_path, 'r') as f:
+            loaded_kwargs = yaml.safe_load(f)
+
+    # Namespace => Dictionary
+    parsed_kwargs = vars(args)
+    parsed_kwargs.update(optional_kwargs)
+
+    kwargs = {}
+    kwargs.update(parsed_kwargs)
+    kwargs.update(loaded_kwargs)
+
+    args = Config(**kwargs)
+
+    # Bind optimizer class.
+    verbose = False
+    args.optimizer = get_optimizer(args.optim, verbose=verbose)
+
+    # Set seeds
+    torch.manual_seed(args.seed)
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+
+    return args
+
+
+class Config(object):
+    def __init__(self, **kwargs):
+        """Configuration Class: set kwargs as class attributes with setattr"""
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    @property
+    def config_str(self):
+        return pprint.pformat(self.__dict__)
+
+    def __repr__(self):
+        """Pretty-print configurations in alphabetical order"""
+        config_str = 'Configurations\n'
+        config_str += self.config_str
+        return config_str
+
+    # def update(self, **kwargs):
+    #     for k, v in kwargs.items():
+    #         setattr(self, k, v)
+
+    # def save(self, path):
+    #     with open(path, 'w') as f:
+    #         yaml.dump(self.__dict__, f, default_flow_style=False)
+
+    # @classmethod
+    # def load(cls, path):
+    #     with open(path, 'r') as f:
+    #         kwargs = yaml.load(f)
+
+    #     return Config(**kwargs)
+
+
+if __name__ == '__main__':
+    args = parse_args(True)

retrieval/pth_loader.py

@@ -0,0 +1,334 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+import h5py
+from lmdbdict import lmdbdict
+from lmdbdict.methods import DUMPS_FUNC, LOADS_FUNC
+import os
+import numpy as np
+import numpy.random as npr
+import random
+
+import torch
+import torch.utils.data as data
+
+import multiprocessing
+import six
+
+verbose = True
+# import torch
+# if torch.cuda.current_device() in [0, -1]:
+if 'LOCAL_RANK' in os.environ and os.environ['LOCAL_RANK'] != '0':
+    verbose = False
+
+class HybridLoader:
+    """
+    If db_path is a director, then use normal file loading
+    If lmdb, then load from lmdb
+    The loading method depend on extention.
+
+    in_memory: if in_memory is True, we save all the features in memory
+               For individual np(y|z)s, we don't need to do that because the system will do this for us.
+               Should be useful for lmdb or h5.
+               (Copied this idea from vilbert)
+    """
+    def __init__(self, db_path, ext, in_memory=False):
+        self.db_path = db_path
+        self.ext = ext
+        if self.ext == '.npy':
+            self.loader = lambda x: np.load(six.BytesIO(x))
+        else:
+            self.loader = lambda x: np.load(six.BytesIO(x))['feat']
+        if db_path.endswith('.lmdb'):
+            self.db_type = 'lmdb'
+            self.lmdb = lmdbdict(db_path, unsafe=True)
+            self.lmdb._key_dumps = DUMPS_FUNC['ascii']
+            self.lmdb._value_loads = LOADS_FUNC['identity']
+        elif db_path.endswith('.pth'): # Assume a key,value dictionary
+            self.db_type = 'pth'
+            self.feat_file = torch.load(db_path)
+            self.loader = lambda x: x
+            print('HybridLoader: ext is ignored')
+        elif db_path.endswith('h5'):
+            self.db_type = 'h5'
+            self.loader = lambda x: np.array(x).astype('float32')
+        else:
+            self.db_type = 'dir'
+
+        self.in_memory = in_memory
+        if self.in_memory:
+            self.features = {}
+
+    def get(self, key):
+
+        if self.in_memory and key in self.features:
+            # We save f_input because we want to save the
+            # compressed bytes to save memory
+            f_input = self.features[key]
+        elif self.db_type == 'lmdb':
+            f_input = self.lmdb[key]
+        elif self.db_type == 'pth':
+            f_input = self.feat_file[key]
+        elif self.db_type == 'h5':
+            f_input = h5py.File(self.db_path, 'r')[key]
+        else:
+            f_input = open(os.path.join(self.db_path, key + self.ext), 'rb').read()
+
+        if self.in_memory and key not in self.features:
+            self.features[key] = f_input
+
+        # load image
+        feat = self.loader(f_input)
+
+        return feat
+
+class CaptionDataset(data.Dataset):
+    
+    def get_vocab_size(self):
+        return self.vocab_size
+
+    def get_vocab(self):
+        return self.ix_to_word
+
+    def get_seq_length(self):
+        return self.seq_length
+
+    def __init__(self, opt):
+        self.opt = opt
+        self.seq_per_img = opt.seq_per_img
+        
+        # feature related options
+        self.use_fc = getattr(opt, 'use_fc', True)
+        self.use_att = getattr(opt, 'use_att', True)
+        self.use_box = getattr(opt, 'use_box', 0)
+        self.norm_att_feat = getattr(opt, 'norm_att_feat', 0)
+        self.norm_box_feat = getattr(opt, 'norm_box_feat', 0)
+
+        # load the json file which contains additional information about the dataset
+        if verbose:
+            print('DataLoader loading json file: ', opt.input_json)
+        self.info = json.load(open(self.opt.input_json))
+        if 'ix_to_word' in self.info:
+            self.ix_to_word = self.info['ix_to_word']
+            self.vocab_size = len(self.ix_to_word)
+            if verbose:
+                print('vocab size is ', self.vocab_size)
+        
+        # open the hdf5 file
+        if verbose:
+            print('DataLoader loading h5 file: ', opt.input_fc_dir, opt.input_att_dir, opt.input_box_dir, opt.input_label_h5)
+        """
+        Setting input_label_h5 to none is used when only doing generation.
+        For example, when you need to test on coco test set.
+        """
+        if self.opt.input_label_h5 != 'none':
+            self.h5_label_file = h5py.File(self.opt.input_label_h5, 'r', driver='core')
+            # load in the sequence data
+            seq_size = self.h5_label_file['labels'].shape
+            self.label = self.h5_label_file['labels'][:]
+            self.seq_length = seq_size[1]
+            if verbose:
+                print('max sequence length in data is', self.seq_length)
+            # load the pointers in full to RAM (should be small enough)
+            self.label_start_ix = self.h5_label_file['label_start_ix'][:]
+            self.label_end_ix = self.h5_label_file['label_end_ix'][:]
+        else:
+            self.seq_length = 1
+
+        self.data_in_memory = getattr(opt, 'data_in_memory', False)
+        self.fc_loader = HybridLoader(self.opt.input_fc_dir, '.npy', in_memory=self.data_in_memory)
+        self.att_loader = HybridLoader(self.opt.input_att_dir, '.npz', in_memory=self.data_in_memory)
+        self.box_loader = HybridLoader(self.opt.input_box_dir, '.npy', in_memory=self.data_in_memory)
+
+        self.use_clipscore = getattr(opt, 'use_clipscore', False)
+        if self.use_clipscore:
+            self.clipscore_loader = HybridLoader(self.opt.input_clipscore_vis_dir, '.npy', in_memory=self.data_in_memory)
+
+
+        self.num_images = len(self.info['images']) # self.label_start_ix.shape[0]
+        if verbose:
+            print('read %d image features' %(self.num_images))
+
+        # separate out indexes for each of the provided splits
+        self.split_ix = {'train': [], 'val': [], 'test': []}
+        for ix in range(len(self.info['images'])):
+            img = self.info['images'][ix]
+            if not 'split' in img:
+                self.split_ix['train'].append(ix)
+                self.split_ix['val'].append(ix)
+                self.split_ix['test'].append(ix)
+            elif img['split'] == 'train':
+                self.split_ix['train'].append(ix)
+            elif img['split'] == 'val':
+                self.split_ix['val'].append(ix)
+            elif img['split'] == 'test':
+                self.split_ix['test'].append(ix)
+            elif opt.train_only == 0: # restval
+                self.split_ix['train'].append(ix)
+
+        if verbose:
+            print('assigned %d images to split train' %len(self.split_ix['train']))
+            print('assigned %d images to split val' %len(self.split_ix['val']))
+            print('assigned %d images to split test' %len(self.split_ix['test']))
+
+    def get_captions(self, ix, seq_per_img):
+        # fetch the sequence labels
+        ix1 = self.label_start_ix[ix] - 1 #label_start_ix starts from 1
+        ix2 = self.label_end_ix[ix] - 1
+        ncap = ix2 - ix1 + 1 # number of captions available for this image
+        assert ncap > 0, 'an image does not have any label. this can be handled but right now isn\'t'
+
+        if ncap < seq_per_img:
+            # we need to subsample (with replacement)
+            seq = np.zeros([seq_per_img, self.seq_length], dtype = 'int')
+            for q in range(seq_per_img):
+                ixl = random.randint(ix1,ix2)
+                seq[q, :] = self.label[ixl, :self.seq_length]
+        else:
+            ixl = random.randint(ix1, ix2 - seq_per_img + 1)
+            seq = self.label[ixl: ixl + seq_per_img, :self.seq_length]
+
+        return seq
+
+    def collate_func(self, batch):
+        seq_per_img = self.seq_per_img
+
+        fc_batch = []
+        att_batch = []
+        label_batch = []
+
+        clip_vis_feat_batch = []
+
+        wrapped = False
+
+        infos = []
+        gts = []
+
+        for sample in batch:
+            # fetch image
+            if self.use_clipscore:
+                tmp_fc, tmp_att, tmp_seq, \
+                    ix, tmp_clip_vis_feat = sample
+
+                clip_vis_feat_batch.append(tmp_clip_vis_feat)
+            else:
+                tmp_fc, tmp_att, tmp_seq, \
+                    ix = sample
+
+            fc_batch.append(tmp_fc)
+            att_batch.append(tmp_att)
+            
+            tmp_label = np.zeros([seq_per_img, self.seq_length + 2], dtype = 'int')
+            if hasattr(self, 'h5_label_file'):
+                # if there is ground truth
+                tmp_label[:, 1 : self.seq_length + 1] = tmp_seq
+            label_batch.append(tmp_label)
+
+            # Used for reward evaluation
+            if hasattr(self, 'h5_label_file'):
+                # if there is ground truth
+                gts.append(self.label[self.label_start_ix[ix] - 1: self.label_end_ix[ix]])
+            else:
+                gts.append([])
+        
+            # record associated info as well
+            info_dict = {}
+            info_dict['ix'] = ix
+            info_dict['id'] = self.info['images'][ix]['id']
+            info_dict['file_path'] = self.info['images'][ix].get('file_path', '')
+            infos.append(info_dict)
+
+        # #sort by att_feat length
+        # fc_batch, att_batch, label_batch, gts, infos = \
+        #     zip(*sorted(zip(fc_batch, att_batch, np.vsplit(label_batch, batch_size), gts, infos), key=lambda x: len(x[1]), reverse=True))
+        if self.use_clipscore:
+            fc_batch, att_batch, label_batch, clip_vis_feat_batch, gts, infos = \
+                zip(*sorted(zip(fc_batch, att_batch, label_batch, clip_vis_feat_batch, gts, infos), key=lambda x: 0, reverse=True))
+        else:
+            fc_batch, att_batch, label_batch, gts, infos = \
+                zip(*sorted(zip(fc_batch, att_batch, label_batch, gts, infos), key=lambda x: 0, reverse=True))
+        data = {}
+        data['fc_feats'] = np.stack(fc_batch)
+        # merge att_feats
+        max_att_len = max([_.shape[0] for _ in att_batch])
+        data['att_feats'] = np.zeros([len(att_batch), max_att_len, att_batch[0].shape[1]], dtype = 'float32')
+        for i in range(len(att_batch)):
+            data['att_feats'][i, :att_batch[i].shape[0]] = att_batch[i]
+        data['att_masks'] = np.zeros(data['att_feats'].shape[:2], dtype='float32')
+        for i in range(len(att_batch)):
+            data['att_masks'][i, :att_batch[i].shape[0]] = 1
+        # set att_masks to None if attention features have same length
+        if data['att_masks'].sum() == data['att_masks'].size:
+            data['att_masks'] = None
+
+        if self.use_clipscore:
+            data['clip_vis_feats'] = np.stack(clip_vis_feat_batch)
+
+        data['labels'] = np.vstack(label_batch)
+        # generate mask
+        nonzeros = np.array(list(map(lambda x: (x != 0).sum()+2, data['labels'])))
+        mask_batch = np.zeros([data['labels'].shape[0], self.seq_length + 2], dtype = 'float32')
+        for ix, row in enumerate(mask_batch):
+            row[:nonzeros[ix]] = 1
+        data['masks'] = mask_batch
+        data['labels'] = data['labels'].reshape(len(batch), seq_per_img, -1)
+        data['masks'] = data['masks'].reshape(len(batch), seq_per_img, -1)
+
+        data['gts'] = gts # all ground truth captions of each images
+        data['infos'] = infos
+
+        data = {k:torch.from_numpy(v) if type(v) is np.ndarray else v for k,v in data.items()} # Turn all ndarray to torch tensor
+
+        return data
+
+    def __getitem__(self, ix):
+        """This function returns a tuple that is further passed to collate_fn
+        """
+        if self.use_att:
+            att_feat = self.att_loader.get(str(self.info['images'][ix]['id']))
+            # Reshape to K x C
+            att_feat = att_feat.reshape(-1, att_feat.shape[-1])
+            if self.norm_att_feat:
+                att_feat = att_feat / np.linalg.norm(att_feat, 2, 1, keepdims=True)
+            if self.use_box:
+                box_feat = self.box_loader.get(str(self.info['images'][ix]['id']))
+                # devided by image width and height
+                x1,y1,x2,y2 = np.hsplit(box_feat, 4)
+                h,w = self.info['images'][ix]['height'], self.info['images'][ix]['width']
+                box_feat = np.hstack((x1/w, y1/h, x2/w, y2/h, (x2-x1)*(y2-y1)/(w*h))) # question? x2-x1+1??
+                if self.norm_box_feat:
+                    box_feat = box_feat / np.linalg.norm(box_feat, 2, 1, keepdims=True)
+                att_feat = np.hstack([att_feat, box_feat])
+                # sort the features by the size of boxes
+                att_feat = np.stack(sorted(att_feat, key=lambda x:x[-1], reverse=True))
+        else:
+            att_feat = np.zeros((0,0), dtype='float32')
+        if self.use_fc:
+            try:
+                fc_feat = self.fc_loader.get(str(self.info['images'][ix]['id']))
+            except:
+                # Use average of attention when there is no fc provided (For bottomup feature)
+                fc_feat = att_feat.mean(0)
+        else:
+            fc_feat = np.zeros((0), dtype='float32')
+        if hasattr(self, 'h5_label_file'):
+            seq = self.get_captions(ix, self.seq_per_img)
+        else:
+            seq = None
+
+        if self.use_clipscore:
+            clip_vis_feat = self.clipscore_loader.get(
+                str(self.info['images'][ix]['id']))
+
+            return (fc_feat,
+                    att_feat, seq,
+                    ix, clip_vis_feat)
+
+        return (fc_feat,
+                att_feat, seq,
+                ix)
+
+    def __len__(self):
+        return len(self.info['images'])

retrieval/text_utils.py

@@ -0,0 +1,74 @@
+import random
+
+def repeat(text, n_max_gram=3, n_max_repeat=3):
+    """repeat n-grams"""
+    tokens = text.split()
+
+    n_gram = random.randint(1, n_max_gram)
+
+    repeat_token_idx = random.randint(0, len(tokens) - n_gram)
+
+    repeated_tokens = tokens[repeat_token_idx:repeat_token_idx+n_gram]
+
+    n_repeat = random.randint(1, n_max_repeat)
+    for _ in range(n_repeat):
+        insert_idx = random.randint(0, len(tokens))
+        tokens = tokens[:insert_idx] + \
+            repeated_tokens + tokens[insert_idx:]
+
+    new_text = " ".join(tokens)
+    return new_text
+
+def remove(text, n_max_gram=3):
+    """remove n-grams"""
+    tokens = text.split()
+
+    n_gram = random.randint(1, n_max_gram)
+
+    remove_token_idx = random.randint(0, len(tokens) - n_gram)
+
+    tokens = tokens[:remove_token_idx] + tokens[remove_token_idx + n_gram:]
+
+    new_text = " ".join(tokens)
+    return new_text
+
+def insert(text, vocab, n_max_tokens=3):
+    """Insert tokens"""
+    tokens = text.split()
+
+    n_insert_token = random.randint(1, n_max_tokens)
+
+    for _ in range(n_insert_token):
+        insert_token_idx = random.randint(0, len(tokens) - 1)
+        insert_token = random.choice(vocab)
+        tokens = tokens[:insert_token_idx] + [insert_token] + tokens[insert_token_idx:]
+
+    new_text = " ".join(tokens)
+    return new_text
+
+def swap(text, vocab, n_max_tokens=3):
+    """Swap tokens"""
+    tokens = text.split()
+
+    n_swap_tokens = random.randint(1, n_max_tokens)
+
+    for _ in range(n_swap_tokens):
+        swap_token_idx = random.randint(0, len(tokens) - 1)
+
+        swap_token = random.choice(vocab)
+        while swap_token == tokens[swap_token_idx]:
+            swap_token = random.choice(vocab)
+
+        tokens[swap_token_idx] = swap_token
+
+    new_text = " ".join(tokens)
+    return new_text
+
+def shuffle(text):
+    """shuffle tokens"""
+    tokens = text.split()
+
+    random.shuffle(tokens)
+
+    new_text = " ".join(tokens)
+    return new_text

retrieval/train_pl.py

@@ -0,0 +1,661 @@
+from ast import parse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+import numpy as np
+
+import time
+import os
+from collections import defaultdict
+
+# import captioning.utils.opts as opts
+# import captioning.models as models
+# from captioning.data.pth_loader import CaptionDataset
+# import captioning.utils.eval_utils as eval_utils
+# import captioning.utils.misc as utils
+# from captioning.utils.rewards import init_scorer, get_self_critical_reward
+# from captioning.modules.loss_wrapper import LossWrapper
+
+from clip_model import CLIPScore
+from caption_data import COCORetrievalDataset
+
+import pytorch_lightning as pl
+
+import detectron2.utils.comm as d2comm
+from detectron2.utils.env import seed_all_rng
+seed_all_rng(1234)
+
+
+class LitModel(pl.LightningModule):
+    def __init__(self, opt):
+        super().__init__()
+        self.opt = opt
+        self.args = args
+        # Intilaize dataset
+        # self.dataset = CaptionDataset(opt)
+
+        # self.dataset = 
+
+        # opt.vocab_size = self.dataset.vocab_size
+        # opt.seq_length = self.dataset.seq_length
+        # self.batch_size = opt.batch_size
+
+        # Build model
+        # opt.vocab = self.dataset.get_vocab()
+        # model = models.setup(opt)
+        # print(model)
+        # del opt.vocab
+
+        # wrapper with loss in it.
+        # lw_model = LossWrapper(model, opt)
+
+        self.model = CLIPScore(use_grammar=opt.use_grammar, joint_out=opt.joint_out)
+        # self.lw_model = lw_model
+
+        for p in self.model.clip_model.vision_model.parameters():
+            p.requires_grad = False
+        for p in self.model.clip_model.visual_projection.parameters():
+            p.requires_grad = False
+
+        # self.struc_flag = None
+        # self.sc_flag = None
+
+
+    def forward(self, *args, **kwargs):
+        """
+        I hate this design. Never pretend it as a nn.Module
+        """
+        raise NotImplementedError
+
+    def train_dataloader(self):
+        # train_dataset = torch.utils.data.Subset(
+        #     self.dataset,
+        #     self.dataset.split_ix['train']
+        # )
+
+        # train_loader = torch.utils.data.DataLoader(
+        #     dataset=train_dataset,
+        #     batch_size=self.batch_size,
+        #     shuffle=True,
+        #     num_workers=4,
+        #     collate_fn=self.dataset.collate_func
+        # )
+
+        train_dataset = COCORetrievalDataset(
+            split='karpathy_train', mode='train',
+            args=opt,
+            verbose=verbose
+            )
+
+        train_loader = torch.utils.data.DataLoader(
+            dataset=train_dataset,
+            batch_size=opt.batch_size,
+            shuffle=True,
+            num_workers=4,
+            collate_fn=train_dataset.collate_fn
+        )
+
+        return train_loader
+
+    def val_dataloader(self, split='karpathy_val'):
+        # val_dataset = torch.utils.data.Subset(
+        #     self.dataset,
+        #     self.dataset.split_ix[split]
+        # )
+        # val_loader = torch.utils.data.DataLoader(
+        #     val_dataset,
+        #     batch_size=self.batch_size,
+        #     shuffle=False,
+        #     num_workers=4,
+        #     drop_last=False,
+        #     collate_fn=self.dataset.collate_func
+        # )
+
+        val_dataset = COCORetrievalDataset(
+            split=split, mode='val',
+            args=opt,
+            verbose=verbose
+        )
+
+        val_loader = torch.utils.data.DataLoader(
+            dataset=val_dataset,
+            batch_size=opt.valid_batch_size,
+            shuffle=False,
+            num_workers=4,
+            drop_last=False,
+            collate_fn=val_dataset.collate_fn
+        )
+
+        return val_loader
+
+    def test_dataloader(self):
+
+        return self.val_dataloader('karpathy_test')
+
+    def training_step(self, data, batch_idx):
+
+
+        batch = data
+        self.model.train()
+
+        model_out = self.model.train_step(
+            img_feat=batch['img_feats'],
+            text=batch['text'],
+            neg_text=batch['neg_text'],
+        )
+
+        clip_loss = model_out['clip_loss']
+
+        if self.opt.joint_out:
+            loss = clip_loss
+        else:
+            grammar_loss = model_out['grammar_loss']
+            loss = clip_loss + grammar_loss
+            
+
+        data_time = self.trainer.profiler.recorded_durations["get_train_batch"][-1]
+        data_time = torch.tensor(data_time)
+
+        # print('batch_idx', batch_idx)
+        # print('loss:', loss)
+
+        # logger_logs = model_out.copy()
+        logger_logs = {}
+        
+        logger_logs['loss'] = loss.detach()
+
+        logger_logs['clip_loss'] = clip_loss.detach()
+
+        if not self.opt.joint_out:
+            logger_logs['grammar_loss'] = grammar_loss.detach()
+
+        logger_logs['data_time'] = data_time.detach()
+
+        # UserWarning: The {progress_bar:dict keyword} was deprecated in 0.9.1 and will be removed in 1.0.0
+        # Please use self.log(...) inside the lightningModule instead.
+
+        # # log on a step or aggregate epoch metric to the logger and/or progress bar
+        # # (inside LightningModule)
+        # self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
+        # warnings.warn(*args, **kwargs)
+        # UserWarning: The {log:dict keyword} was deprecated in 0.9.1 and will be removed in 1.0.0
+        # Please use self.log(...) inside the lightningModule instead.
+
+        # output = {
+        #     'loss': loss,
+        #     'log': logger_logs,
+        #     'progress_bar': {'data_time': data_time}
+        # }
+
+        for k, v in logger_logs.items():
+            if k in ['data_time', 'clip_loss', 'grammar_loss']:
+                self.log('train/'+k, v, prog_bar=True)
+            else:
+                self.log('train/'+k, v)
+        
+        # print('training step logged')
+
+        return loss
+
+    def validation_step(self, data, batch_idx):
+
+        batch = data
+        self.model.eval()
+
+        with torch.no_grad():
+            model_out = self.model.train_step(
+                img_feat=batch['img_feats'],
+                text=batch['text'],
+                neg_text=batch['neg_text'],
+            )
+
+            if self.opt.joint_out:
+                clip_loss = model_out['clip_loss']
+                loss = clip_loss
+
+                output = {
+                    # 'val_loss': loss,
+                    'loss': loss.detach(),
+                    'clip_loss': clip_loss.detach(),
+                    # 'grammar_loss': grammar_loss.detach(),
+
+                    'img_feat': model_out['img_feat'].detach(),
+                    'text_feat': model_out['text_feat'].detach(),
+                    # 'neg_text_feat': model_out['neg_text_feat'].detach(),
+                    # 'grammar_pos_pred': model_out['grammar_pos_pred'].detach(),
+                    # 'grammar_neg_pred': model_out['grammar_neg_pred'].detach(),
+                    # 'predictions': predictions,
+                    # 'n_predictions': n_predictions,
+                }
+            else:
+                clip_loss = model_out['clip_loss']
+                grammar_loss = model_out['grammar_loss']
+                loss = clip_loss + grammar_loss
+
+                output = {
+                    # 'val_loss': loss,
+                    'loss': loss.detach(),
+                    'clip_loss': clip_loss.detach(),
+                    'grammar_loss': grammar_loss.detach(),
+
+                    'img_feat': model_out['img_feat'].detach(),
+                    'text_feat': model_out['text_feat'].detach(),
+                    # 'neg_text_feat': model_out['neg_text_feat'].detach(),
+                    'grammar_pos_pred': model_out['grammar_pos_pred'].detach(),
+                    'grammar_neg_pred': model_out['grammar_neg_pred'].detach(),
+                    # 'predictions': predictions,
+                    # 'n_predictions': n_predictions,
+                }
+        return output
+
+    def test_step(self, *args, **kwargs):
+        return self.validation_step(*args, **kwargs)
+
+    def validation_epoch_end(self, outputs, split='val'):
+        outputs = d2comm.gather(outputs)
+        # master node
+        if d2comm.is_main_process():
+            assert self.trainer.node_rank == 0 and self.trainer.local_rank == 0
+            outputs = sum(outputs, [])
+
+            out = {}
+
+            val_loss_mean = sum([_['loss'].cpu() for _ in outputs]) / len(outputs)
+            val_clip_loss_mean = sum([_['clip_loss'].cpu() for _ in outputs]) / len(outputs)
+            if not self.opt.joint_out:
+                val_grammar_loss_mean = sum([_['grammar_loss'].cpu() for _ in outputs]) / len(outputs)
+
+            print('loss', val_loss_mean.item())
+            print('clip_loss', val_clip_loss_mean.item())
+            if not self.opt.joint_out:
+                print('grammar_loss', val_grammar_loss_mean.item())
+
+            logit_scale = self.model.clip_model.logit_scale.exp().cpu()
+
+            text_feats = torch.cat([_['text_feat'].cpu() for _ in outputs], dim=0)
+            img_feats = torch.cat([_['img_feat'].cpu() for _ in outputs], dim=0)
+
+            assert text_feats.size() == (5000, 512), text_feats.size()
+            assert img_feats.size() == (5000, 512), img_feats.size()
+
+            logits_per_text = torch.matmul(text_feats, img_feats.t()) * logit_scale
+            logits_per_image = logits_per_text.T
+
+            # text-to-image retrieval
+            print('Text-to-Image retrieval')
+            for k in [1, 5, 10]:
+                text_to_image_topk = logits_per_text.topk(k, dim=1).indices
+
+                n_text = len(text_to_image_topk)
+
+                labels = torch.arange(0, n_text).view(-1, 1)
+
+                n_retrieved = ((text_to_image_topk == labels).sum(dim=1) > 0).sum()
+
+                recall_k = n_retrieved / n_text * 100
+
+                out[f'text_to_image_recall_{k}'] = recall_k.item()
+
+                print(f'R@{k}: {recall_k.item():.2f}%')
+
+            # image-to-text retrieval
+            print('Image-to-Text retrieval')
+            for k in [1, 5, 10]:
+                image_to_text_topk = logits_per_image.topk(k, dim=1).indices
+
+                n_image = len(image_to_text_topk)
+
+                labels = torch.arange(0, n_image).view(-1, 1)
+
+                n_retrieved = ((image_to_text_topk == labels).sum(dim=1) > 0).sum()
+
+                recall_k = n_retrieved / n_image * 100
+
+                out[f'image_to_text_recall_{k}'] = recall_k.item()
+
+                print(f'R@{k}: {recall_k.item():.2f}%')
+
+            out.update({
+                'loss': val_loss_mean.item(),
+                'clip_loss': val_clip_loss_mean.item()
+            })
+
+            if not self.opt.joint_out:
+                # grammar scoring
+                grammar_pos_pred = torch.cat([_['grammar_pos_pred'].cpu() for _ in outputs], dim=0)
+                grammar_neg_pred = torch.cat([_['grammar_neg_pred'].cpu() for _ in outputs], dim=0)
+
+                TP = (grammar_pos_pred == 1).sum().item()
+                FP = (grammar_pos_pred == 0).sum().item()
+                FN = (grammar_neg_pred == 1).sum().item()
+                TN = (grammar_neg_pred == 0).sum().item()
+                print('Grammar check')
+                print(f'TP: {TP} FP: {FP}  FN: {FN}  TN: {TN}')
+
+                precision = TP / (TP + FP) * 100
+                recall = TP / (TP + FN) * 100
+                accuracy = (TP + TN) / (TP + FP + FN + TN) * 100
+                f1 = 2 * precision * recall / (precision + recall)
+                print(f'Precision: {precision:.2f}%')
+                print(f'Recall: {recall:.2f}%')
+                print(f'Accuracy: {accuracy:.2f}%')
+                print(f'F1: {f1:.2f}%')
+                print('Total: {}'.format(len(grammar_pos_pred)))
+
+                out.update({
+                    'grammar_loss': val_grammar_loss_mean,
+
+                    'grammar_precision': precision,
+                    'grammar_recall': recall,
+                    'grammar_accuracy': accuracy,
+                    'grammar_f1': f1,
+
+                })
+
+        else:
+            out = {}
+
+        out = d2comm.all_gather(out)[0]  # Only the one from master node
+        assert len(out) > 0  # make sure the head has index 0
+
+        # must all be tensors
+        out = {k: torch.tensor(v) if not torch.is_tensor(
+            v) else v for k, v in out.items()}
+
+        for k, v in out.items():
+            self.log(f'{split}/{k}', v)
+
+    def test_epoch_end(self, outputs):
+
+        self.validation_epoch_end(outputs, 'test')
+        
+    def configure_optimizers(self):
+        # opt = self.opt
+        # model = self.model
+
+        # parameters = [p for p in model.parameters() if p.requires_grad]
+
+        # if opt.noamopt:
+        #     # assert opt.caption_model in ['transformer', 'bert', 'm2transformer'], 'noamopt can only work with transformer'
+        #     optimizer = utils.get_std_opt(
+        #         model, optim_func=opt.optim, factor=opt.noamopt_factor, warmup=opt.noamopt_warmup)
+        # elif opt.reduce_on_plateau:
+        #     # optimizer = utils.build_optimizer(model.parameters(), opt)
+        #     optimizer = utils.build_optimizer(parameters, opt)
+        #     optimizer = utils.ReduceLROnPlateau(optimizer,
+        #                                         factor=opt.reduce_on_plateau_factor,
+        #                                         patience=opt.reduce_on_plateau_patience)
+        # else:
+        #     # optimizer = utils.build_optimizer(model.parameters(), opt)
+        #     optimizer = utils.build_optimizer(parameters, opt)
+
+
+        # from transformers.optimization import AdamW, get_linear_schedule_with_warmup
+        # batch_per_epoch = len(self.train_loader)
+        # t_total = batch_per_epoch // self.args.gradient_accumulation_steps * self.args.epochs
+        # warmup_ratio = self.args.warmup_ratio
+        # warmup_iters = int(t_total * warmup_ratio)
+        # if self.verbose:
+        #     print("Batch per epoch: %d" % batch_per_epoch)
+        #     print("Total Iters: %d" % t_total)
+        #     print('Warmup ratio:', warmup_ratio)
+        #     print("Warm up Iters: %d" % warmup_iters)
+
+        if self.args.optim == 'adamw':
+            no_decay = ["bias", "LayerNorm.weight"]
+            optimizer_grouped_parameters = [
+                {
+                    "params": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay)],
+                    "weight_decay": self.args.weight_decay,
+                },
+                {
+                    "params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay)],
+                    "weight_decay": 0.0,
+                },
+            ]
+
+            for group in optimizer_grouped_parameters:
+                group['params'] = [p for p in group['params'] if p.requires_grad]
+
+            from transformers.optimization import AdamW
+            optim = AdamW(optimizer_grouped_parameters,
+                            lr=self.args.lr, eps=self.args.adam_eps)
+            # lr_scheduler = get_linear_schedule_with_warmup(
+            #     optim, warmup_iters, t_total)
+
+        # optimizers = []
+        optimizers = [optim]
+        lr_schedulers = []
+
+        return optimizers, lr_schedulers
+
+    def optimizer_step(self, epoch, batch_idx, optimizer,
+                       optimizer_idx, *args, **kwargs):
+        # # warm up lr
+        # opt = self.opt
+        # iteration = self.trainer.global_step
+        # if opt.use_warmup and (iteration < opt.noamopt_warmup):
+        #     opt.current_lr = opt.learning_rate * \
+        #         (iteration+1) / opt.noamopt_warmup
+        #     utils.set_lr(optimizer, opt.current_lr)
+
+        super().optimizer_step(epoch, batch_idx, optimizer,
+                               optimizer_idx, *args, **kwargs)
+
+        # print('optimizer step')
+
+    def state_dict(self):
+        """
+        Save the model state dict as well as opt and vocab
+        """
+        state_dict = self.model.state_dict()
+        device = next(iter(state_dict.values())).device
+        assert '_vocab' not in state_dict and '_opt' not in state_dict, 'Just in case'
+        # state_dict.update({
+        #     '_vocab': utils.serialize_to_tensor(self.model.vocab).to(device),
+        #     '_opt': utils.serialize_to_tensor(self.opt).to(device)
+        # })
+        return state_dict
+
+    def load_state_dict(self, state_dict=None, strict=True):
+        # if '_vocab' in state_dict:
+        #     self.model.vocab = utils.deserialize(state_dict['_vocab'])
+        #     del state_dict['_vocab']
+        # elif strict:
+        #     raise KeyError
+        # if '_opt' in state_dict:
+        #     saved_model_opt = utils.deserialize(state_dict['_opt'])
+        #     del state_dict['_opt']
+        #     opt = self.opt
+        #     # Make sure the saved opt is compatible with the curren topt
+        #     need_be_same = ["caption_model",
+        #                     "rnn_type", "rnn_size", "num_layers"]
+        #     for checkme in need_be_same:
+        #         if getattr(saved_model_opt, checkme) in ['updown', 'topdown'] and \
+        #                 getattr(opt, checkme) in ['updown', 'topdown']:
+        #             continue
+        #         assert getattr(saved_model_opt, checkme) == getattr(
+        #             opt, checkme), "Command line argument and saved model disagree on '%s' " % checkme
+        # elif strict:
+        #     raise KeyError
+        self.model.load_state_dict(state_dict, strict)
+
+
+class OnEpochStartCallback(pl.Callback):
+
+    def on_epoch_start(self, trainer, pl_module):
+        # Update lr/training stage/scheduled sampling prob etc.
+        opt = pl_module.opt
+        model = pl_module.model
+        epoch = trainer.current_epoch
+        optimizer = trainer.optimizers[0]
+
+        # if not opt.noamopt and not opt.reduce_on_plateau:
+        #     # Assign the learning rate
+        #     if epoch > opt.learning_rate_decay_start and opt.learning_rate_decay_start >= 0:
+        #         frac = (
+        #             epoch - opt.learning_rate_decay_start) // opt.learning_rate_decay_every
+        #         decay_factor = opt.learning_rate_decay_rate ** frac
+        #         opt.current_lr = opt.learning_rate * decay_factor
+        #     else:
+        #         opt.current_lr = opt.learning_rate
+        #     utils.set_lr(optimizer, opt.current_lr)  # set the decayed rate
+        # # Assign the scheduled sampling prob
+        # if epoch > opt.scheduled_sampling_start and opt.scheduled_sampling_start >= 0:
+        #     frac = (
+        #         epoch - opt.scheduled_sampling_start) // opt.scheduled_sampling_increase_every
+        #     opt.ss_prob = min(opt.scheduled_sampling_increase_prob *
+        #                       frac, opt.scheduled_sampling_max_prob)
+        #     model.ss_prob = opt.ss_prob
+
+        # # If start self critical training
+        # if opt.self_critical_after != -1 and epoch >= opt.self_critical_after:
+        #     sc_flag = True
+        #     init_scorer(opt.cached_tokens)
+        # else:
+        #     sc_flag = False
+
+        # # If start structure loss training
+        # if opt.structure_after != -1 and epoch >= opt.structure_after:
+        #     struc_flag = True
+        #     init_scorer(opt.cached_tokens)
+        # else:
+        #     struc_flag = False
+
+        # pl_module.struc_flag = struc_flag
+        # pl_module.sc_flag = sc_flag
+
+
+class ModelCheckpoint(pl.callbacks.ModelCheckpoint):
+
+    def on_keyboard_interrupt(self, trainer, pl_module):
+        # Save model when keyboard interrupt
+        filepath = os.path.join(self.dirpath, self.prefix + 'interrupt.ckpt')
+        self._save_model(filepath)
+
+from param import parse_args
+# opt = opts.parse_opt()
+args = parse_args()
+opt = args
+
+checkpoint_callback = ModelCheckpoint(
+    filepath=opt.checkpoint_dir + '{epoch:02d}',
+    # dirpath=opt.checkpoint_path,
+    save_last=True,
+    save_top_k=1,
+    verbose=True,
+    # monitor='to_monitor',
+    # monitor='val/to_monitor',
+    # monitor='val/CIDEr',
+    monitor='val/loss',
+    mode='min',
+    # prefix=opt.id+'_',
+    prefix=opt.id,
+    # filename=f'{opt.id}_',
+)
+
+verbose = True
+# import torch
+# if torch.cuda.current_device() in [0, -1]:
+if 'LOCAL_RANK' in os.environ and os.environ['LOCAL_RANK'] != '0':
+    verbose = False
+
+# if verbose:
+#     print(opt)
+#     print("""
+#     val_image_use,
+#     save_checkpoint_very
+#     save_every_epoch,
+#     save_history-ckpt will be ignored.
+#     """)
+
+# Lightning defines batch size as batch size per gpu
+assert opt.batch_size % torch.cuda.device_count() == 0
+opt.batch_size = opt.batch_size // torch.cuda.device_count()
+opt.valid_batch_size = opt.valid_batch_size // torch.cuda.device_count()
+
+# If resume from last checkpoint
+# if opt.start_from is not None and os.path.isfile(os.path.join(opt.start_from, f'{opt.id}_last.ckpt')):
+#     resume_from = os.path.join(opt.start_from, f'{opt.id}_last.ckpt')
+if opt.start_from is not None and os.path.isfile(os.path.join(opt.start_from, f'{opt.id}-last.ckpt')):
+    resume_from = os.path.join(opt.start_from, f'{opt.id}-last.ckpt')
+    if verbose:
+        print('resume from', resume_from)
+else:
+    resume_from = None
+
+from pytorch_lightning.loggers import WandbLogger
+wandb_logger = WandbLogger(
+    # project='CLIP-ViL-COCOCaption',
+    project='CLIP-Finetune-COCO',
+    name=opt.id,
+)
+
+if verbose:
+    wandb_logger.experiment.config.update(opt)
+    from pathlib import Path
+    import glob
+    import wandb
+    # src_dir = Path(__file__).resolve().parent.parent
+    glob_str = "*.py"
+    base_path = './'
+    wandb.save(glob_str=glob_str, base_path=base_path)
+
+    glob_str = "**/*.yaml"
+    base_path = './'
+    wandb.save(glob_str=glob_str, base_path=base_path)
+    
+    # code = wandb.Artifact('project-source', type='code')
+    # for path in glob.glob('**/*.py', recursive=True):
+    #     code.add_file(path, name='source/'+path)
+    #     print(path)
+    # wandb.run.use_artifact(code)
+
+
+
+
+lit = LitModel(opt)
+# warning grad_clip_mode is ignored.
+trainer = pl.Trainer(
+    callbacks=[
+        OnEpochStartCallback(),
+        # pl.callbacks.lr_logger.LearningRateLogger()
+        pl.callbacks.LearningRateMonitor()
+    ],
+    default_root_dir=opt.checkpoint_dir,
+    resume_from_checkpoint=resume_from,
+
+    distributed_backend='ddp',
+    gpus=torch.cuda.device_count(),
+    
+    # gpus=1,
+
+    check_val_every_n_epoch=1,
+    # max_epochs=opt.max_epochs,
+    max_epochs=opt.epochs,
+    # gradient_clip_val=opt.grad_clip_value,
+    gradient_clip_val=opt.clip_grad_norm,
+    
+    checkpoint_callback=checkpoint_callback,
+    log_gpu_memory='min_max',
+    # log_save_interval=opt.losses_log_every,
+    log_every_n_steps=opt.losses_log_every,
+    profiler=True,
+    # profiler='simple',
+    # row_log_interval=10,  # what is it?
+    flush_logs_every_n_steps=10,
+    num_sanity_val_steps=0,
+    # val_check_interval=0.01,
+    # limit_train_batches=500,
+    # progress_bar_refresh_rate=0,
+    # fast_dev_run=True,
+    precision=opt.precision,
+    logger=wandb_logger
+)
+
+if os.getenv('EVALUATE', '0') == '1':
+    trainer.test(lit)
+else:
+    trainer.fit(lit)

save/README.md

@@ -0,0 +1 @@
+Directory for checkpoints