update the operator.

Signed-off-by: wxywb <xy.wang@zilliz.com>

albef.py

@@ -18,19 +18,25 @@ from PIL import Image
 import torch
 import yaml
 from torchvision import transforms
+from urllib.parse import urlparse
+from timm.models.hub import download_cached_file
 
 from towhee.types.image_utils import to_pil
 from towhee.operator.base import NNOperator, OperatorFlag
 from towhee.types.arg import arg, to_image_color
 from towhee import register
 
+def is_url(url_or_filename):
+    parsed = urlparse(url_or_filename)
+    return parsed.scheme in ("http", "https")
+
 @register(output_schema=['vec'])
 class Albef(NNOperator):
     """
     ALBEF multi-modal embedding operator
     """
-    def prepare_model(checkpoint_path, model):
-        checkpoint = torch.load(checkpoint_path, map_location='cpu')
+    def prepare_model(self, checkpoint_path, model, interpolate_pos_embed):
+        checkpoint = self.load_checkpoint(checkpoint_path)
         state_dict = checkpoint['model']
         pos_embed_reshaped = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],model.visual_encoder)
         state_dict['visual_encoder.pos_embed'] = pos_embed_reshaped
@@ -42,23 +48,44 @@ class Albef(NNOperator):
                 state_dict[encoder_key] = state_dict[key]
                 del state_dict[key]
         msg = model.load_state_dict(state_dict,strict=False)
-        print('load checkpoint from ' + checkpoint_path)
         return model
+    
+    def load_checkpoint(self, url_or_filename):
+        if is_url(url_or_filename):
+            cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
+            checkpoint = torch.load(cached_file, map_location='cpu')
+        elif os.path.isfile(url_or_filename):
+            checkpoint = torch.load(url_or_filename, map_location='cpu')
+        else:
+            raise RuntimeError('checkpoint url or path is invalid')
+
+        return checkpoint
 
     def __init__(self, model_name: str, modality: str):
         self.modality = modality
         config = self._configs()[model_name]
 
+        path = str(Path(__file__).parent)
+        sys.path.append(path)
+
+        from models.model_retrieval import ALBEF
+        from models.vit import interpolate_pos_embed
+        from models.tokenization_bert import BertTokenizer
+        sys.path.pop()
+
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
         normalize = transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
-        tokenizer = BertTokenizer.from_pretrained(config)
-        model = ALBEF(config=config, text_encoder=config['text_encoder'], tokenizer=tokenizer)
-        cfg = yaml.load(open(config['cfg'], 'r'), Loader=yaml.Loader)
-        checkpoint_path = cfg['ckpt_path'] 
+        tokenizer = BertTokenizer.from_pretrained(config['text_encoder'])
+        self.tokenizer = tokenizer
 
-        self.model = self.prepare_model(checkpoint_path, model)
+        cfg = yaml.load(open(path + '/' + config['cfg_path'], 'r'), Loader=yaml.Loader)
+        cfg['bert_config'] = path + '/' + cfg['bert_config']
+        model = ALBEF(config=cfg, text_encoder=config['text_encoder'], tokenizer=tokenizer)
+        checkpoint_path = config['weights'] 
 
+        self.model = self.prepare_model(checkpoint_path, model, interpolate_pos_embed)
+        
         self.test_transform = transforms.Compose([
             transforms.Resize((cfg['image_res'],cfg['image_res']),interpolation=Image.BICUBIC),
             transforms.ToTensor(),
@@ -90,16 +117,20 @@ class Albef(NNOperator):
             return results
 
     def _inference_from_text(self, text):
-        tokens = self.text_tokenizer(text, return_tensors='pt', padding=True)['input_ids'].to(self.device)
-        text_features = self.text_encoder(tokens).logits
-        return text_features
+        text_input = self.tokenizer(text, padding='max_length', truncation=True, max_length=30, return_tensors="pt").to(self.device)
+
+        text_output = self.model.text_encoder(text_input.input_ids, attention_mask = text_input.attention_mask, mode='text')
+        text_feat = text_output.last_hidden_state
+        text_embed = self.model.text_proj(text_feat[:,0,:])
+        return text_embed
 
     @arg(1, to_image_color('RGB'))
     def _inference_from_image(self, img):
         image = to_pil(img)
-        image = self.processor(images=image, return_tensors="pt").to(self.device)
-        image_features = self.clip_model.get_image_features(**image)
-        return image_features
+        image = self.test_transform(image).unsqueeze(0)
+        image_feat = self.model.visual_encoder(image)
+        image_embed = self.model.vision_proj(image_feat[:,0,:])
+        return image_embed
 
     def _configs(self):
         config = {}
@@ -107,7 +138,11 @@ class Albef(NNOperator):
         config['albef_4m']['tokenizer'] = 'bert-base-uncased'
         config['albef_4m']['text_encoder'] = 'bert-base-uncased'
         config['albef_4m']['cfg_path'] = './configs/Retrieval_flickr.yaml'
-        config['albef_4m']['ckpt_path'] = ''
-
-
+        config['albef_4m']['weights'] = 'https://storage.googleapis.com/sfr-pcl-data-research/ALBEF/ALBEF_4M.pth'
 
+        config['albef_14m'] = {}
+        config['albef_14m']['tokenizer'] = 'bert-base-uncased'
+        config['albef_14m']['text_encoder'] = 'bert-base-uncased'
+        config['albef_14m']['cfg_path'] = './configs/Retrieval_flickr.yaml'
+        config['albef_14m']['weights'] = 'https://storage.googleapis.com/sfr-pcl-data-research/ALBEF/ALBEF.pth'
+        return config

configs/Grounding.yaml

@@ -0,0 +1,33 @@
+train_file: ['data/refcoco+_train.json']
+test_file: ['data/refcoco+_val.json','data/refcoco+_test.json']
+
+refcoco_data: 'data'
+det_file: 'data/refcoco+/dets.json'            
+coco_file: 'data/refcoco+/cocos.json'    
+
+image_root: '/export/share/datasets/vision/coco/images/'
+
+bert_config: 'configs/config_bert.json'
+
+image_res: 384
+batch_size: 32
+
+queue_size: 65536
+momentum: 0.995
+vision_width: 768
+embed_dim: 256
+temp: 0.07
+
+alpha: 0.4
+distill: True
+warm_up: True
+
+optimizer: {opt: adamW, lr: 1e-5, weight_decay: 0.02}
+schedular: {sched: cosine, lr: 1e-5, epochs: 5, min_lr: 1e-6, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 1, cooldown_epochs: 0}
+
+
+
+
+
+
+

configs/NLVR.yaml

@@ -0,0 +1,25 @@
+train_file: ['data/nlvr_train.json']
+val_file: ['data/nlvr_dev.json']                  
+test_file: ['data/nlvr_test.json']  
+
+image_root: '/export/share/datasets/vision/NLVR2/' 
+
+image_res: 384
+batch_size: 16 
+
+bert_config: 'configs/config_bert.json'
+
+alpha: 0.4
+distill: True
+warm_up: True
+eval_ema: False
+
+optimizer: {opt: adamW, lr: 2e-5, weight_decay: 0.02}
+schedular: {sched: cosine, lr: 2e-5, epochs: 10, min_lr: 1e-6, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 1, cooldown_epochs: 0}
+
+
+
+
+
+
+

configs/NLVR_pretrain.yaml

@@ -0,0 +1,25 @@
+train_file: ['data/coco.json',
+               'data/vg.json',
+               'data/cc3m_train.json',
+               'data/cc3m_val.json',
+               'data/sbu.json'               
+               ]
+
+# each train_file (json) contains a python list where each item is {'image': img_path, 'caption': text or list_of_text }  
+
+bert_config: 'configs/config_bert.json'
+
+image_res: 256
+vision_width: 768
+embed_dim: 256
+batch_size: 64
+
+optimizer: {opt: adamW, lr: 2e-5, weight_decay: 0.02}
+schedular: {sched: cosine, lr: 2e-5, epochs: 1, min_lr: 1e-5, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 1, cooldown_epochs: 0}
+
+
+
+
+
+
+

configs/Pretrain.yaml

@@ -0,0 +1,29 @@
+train_file: ['data/coco.json',
+               'data/vg.json',
+               'data/cc12m.json',
+               'data/cc3m_train.json',
+               'data/cc3m_val.json',
+               'data/sbu.json'               
+               ]
+# each train_file (json) contains a python list where each item is {'image': img_path, 'caption': text or list_of_text }               
+bert_config: 'configs/config_bert.json'
+
+image_res: 256
+vision_width: 768
+embed_dim: 256
+batch_size: 64
+temp: 0.07
+mlm_probability: 0.15
+queue_size: 65536
+momentum: 0.995
+alpha: 0.4
+
+optimizer: {opt: adamW, lr: 1e-4, weight_decay: 0.02}
+schedular: {sched: cosine, lr: 1e-4, epochs: 30, min_lr: 1e-5, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 20, cooldown_epochs: 0}
+
+
+
+
+
+
+

configs/Retrieval_coco.yaml

@@ -0,0 +1,31 @@
+train_file:  ['data/coco_train.json']
+val_file: 'data/coco_val.json'                
+test_file: 'data/coco_test.json'
+image_root: '/export/share/datasets/vision/coco/images/'
+
+bert_config: 'configs/config_bert.json'
+
+image_res: 384
+batch_size_train: 32
+batch_size_test: 64
+
+queue_size: 65536
+momentum: 0.995
+vision_width: 768
+embed_dim: 256
+temp: 0.07
+k_test: 256
+
+alpha: 0.4
+distill: True
+warm_up: True
+
+optimizer: {opt: adamW, lr: 1e-5, weight_decay: 0.02}
+schedular: {sched: cosine, lr: 1e-5, epochs: 5, min_lr: 1e-6, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 1, cooldown_epochs: 0}
+
+
+
+
+
+
+

configs/Retrieval_flickr.yaml

@@ -0,0 +1,31 @@
+train_file:  ['data/flickr30k_train.json']
+val_file: 'data/flickr30k_val.json'                
+test_file: 'data/flickr30k_test.json'
+image_root: '/export/share/datasets/vision/flickr30k/' #flickr30k-images/
+
+bert_config: 'configs/config_bert.json'
+
+image_res: 384
+batch_size_train: 32
+batch_size_test: 64
+
+queue_size: 65536
+momentum: 0.995
+vision_width: 768
+embed_dim: 256
+temp: 0.07
+k_test: 128
+
+alpha: 0.4
+distill: True
+warm_up: True
+
+optimizer: {opt: adamW, lr: 1e-5, weight_decay: 0.02} 
+schedular: {sched: cosine, lr: 1e-5, epochs: 10, min_lr: 1e-6, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 1, cooldown_epochs: 0}
+
+
+
+
+
+
+

configs/VE.yaml

@@ -0,0 +1,25 @@
+train_file: 'data/ve_train.json'
+val_file: 'data/ve_dev.json'                
+test_file: 'data/ve_test.json'
+
+image_root: '/export/home/project/SNLI-VE/data/images'
+
+image_res: 384
+batch_size_train: 32 
+batch_size_test: 64 
+
+alpha: 0.4
+distill: True
+warm_up: False
+
+bert_config: 'configs/config_bert.json'
+
+optimizer: {opt: adamW, lr: 2e-5, weight_decay: 0.02}
+schedular: {sched: cosine, lr: 2e-5, epochs: 5, min_lr: 1e-6, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 1, cooldown_epochs: 0}
+
+
+
+
+
+
+

configs/VQA.yaml

@@ -0,0 +1,32 @@
+train_file: ['data/vqa_train.json',
+             'data/vqa_val.json',   
+             'data/vg_qa.json']
+             
+test_file: ['data/vqa_test.json']
+answer_list: 'data/answer_list.json'
+
+vqa_root: '/export/share/datasets/vision/VQA/Images/mscoco/' #train2014/
+vg_root: '/export/share/datasets/vision/visual-genome/'  #image/
+
+image_res: 384
+batch_size_train: 32 
+batch_size_test: 16 
+k_test: 128
+
+alpha: 0.4
+distill: True
+warm_up: True
+
+eos: '[SEP]'
+
+bert_config: 'configs/config_bert.json'
+
+optimizer: {opt: adamW, lr: 2e-5, weight_decay: 0.02}
+schedular: {sched: cosine, lr: 2e-5, epochs: 8, min_lr: 1e-6, decay_rate: 1, warmup_lr: 1e-5, warmup_epochs: 4, cooldown_epochs: 0}
+
+
+
+
+
+
+

configs/config_bert.json

@@ -0,0 +1,21 @@
+{
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "type_vocab_size": 2,
+  "vocab_size": 30522,
+  "fusion_layer": 6,
+  "encoder_width": 768
+}