Upload 6 files

xdecoder/language/build.py

@@ -0,0 +1,11 @@
+from .registry import model_entrypoints
+from .registry import is_model
+
+
+def build_language_encoder(config, **kwargs):
+    model_name = config['MODEL']['TEXT']['ARCH']
+
+    if not is_model(model_name):
+        raise ValueError(f'Unkown model: {model_name}')
+
+    return model_entrypoints(model_name)(config, **kwargs)

xdecoder/language/fixvlpencoder.py

@@ -0,0 +1,35 @@
+from importlib.metadata import requires
+import torch
+import torch.nn as nn
+
+from .registry import register_model
+from .vlpencoder import LanguageEncoder
+
+class FixLanguageEncoder(LanguageEncoder):
+
+    def __init__(
+        self,
+        *args, **kwargs):
+        super(FixLanguageEncoder, self).__init__(*args, **kwargs)
+        self.logit_scale = nn.Parameter(torch.ones([]), requires_grad=False)
+
+    @torch.no_grad()
+    def get_text_embeddings(self, *args, **kwargs):
+        return super().get_text_embeddings(*args, **kwargs)
+
+    @torch.no_grad()
+    def get_text_token_embeddings(self, *args, **kwargs):
+        return super().get_text_token_embeddings(*args, **kwargs)
+
+    @torch.no_grad()
+    def forward_language(self, *args, **kwargs):
+        return super().forward_language(*args, **kwargs)
+
+    @torch.no_grad()
+    def forward_language_token(self, *args, **kwargs):
+        return super().forward_language_token(*args, **kwargs)
+
+
+@register_model
+def get_language_model(cfg, **kwargs):
+    return FixLanguageEncoder(cfg)        

xdecoder/language/loss.py

@@ -0,0 +1,225 @@
+import pickle
+from distutils import log
+
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from einops import rearrange, repeat
+from timm.loss import SoftTargetCrossEntropy
+
+soft_cross_entropy = SoftTargetCrossEntropy()
+
+def is_dist_initialized():
+    return torch.distributed.is_initialized()
+
+def get_world_size():
+    if is_dist_initialized():
+        return torch.distributed.get_world_size()
+    return 1
+
+def get_rank():
+    if is_dist_initialized():
+        return dist.get_rank()
+    return 0
+
+def all_gather_grad(x):
+    if get_world_size() > 1:
+        all_x = [torch.zeros_like(x) for _ in range(get_world_size())]
+        torch.distributed.all_gather(all_x, x)
+        all_x[torch.distributed.get_rank()] = x
+        x = torch.cat(all_x, dim=0)
+    return x
+
+def vl_multilabel_contrastive_loss(image_feat, text_feat, temperature=1):
+    """
+    Args:
+        image_feat (torch.Tensor): shape [B, L1, C] # B: batch_size, L1: 1, C: 256
+        text_feat (torch.Tensor): shape [B, L2, C] # B:batch_size, L2: number of selected nouns, C: 256
+
+    Returns:
+    """
+    # [B, L1, C], L1 = 1
+    # image_feat = F.normalize(image_feat, dim=-1)
+    # [B, L2, C]
+    # text_feat = F.normalize(text_feat, dim=-1)
+    # HACK: normalize outside
+    
+    # [B, L1, L2]
+    dist_per_img = image_feat @ rearrange(text_feat, 'b l c -> b c l')    
+    # [B, L2, L1]
+    dist_per_text = text_feat @ rearrange(image_feat, 'b l c -> b c l')
+
+    batch = image_feat.shape[0]
+    img_len = image_feat.shape[1]
+    text_len = text_feat.shape[1]
+    # [B, L1, L2]
+    pos_labels_batch_img = rearrange(torch.ones_like(dist_per_text) / dist_per_text.size(1), 'b l2 l1 -> b l1 l2')
+    # [B, L2, L1]
+    pos_labels_batch_text = rearrange(torch.ones_like(dist_per_img) / dist_per_img.size(1), 'b l1 l2 -> b l2 l1')
+
+    image_x = rearrange(image_feat, 'b l c -> (b l) c')
+    text_x = rearrange(text_feat, 'b l c -> (b l) c')
+
+    logits_per_img = image_x @ all_gather_grad(text_x).t()
+    logits_per_text = text_x @ all_gather_grad(image_x).t()
+
+    # get label globally
+    # [B, L1, B, L2, W]
+    labels_per_img = F.one_hot(
+        torch.ones(batch, img_len, batch, text_len, dtype=torch.long, device=image_x.device) * get_rank(),
+        num_classes=get_world_size()).to(image_x.dtype)
+    labels_per_img *= rearrange(pos_labels_batch_img, 'b l1 l2 -> b l1 1 l2 1') * repeat(
+        torch.eye(batch, dtype=image_x.dtype, device=image_x.device), 'b1 b2 -> b1 1 b2 1 1')
+    # [BxL1, WxBxL2]
+    labels_per_img = rearrange(labels_per_img, 'b1 l1 b2 l2 w -> (b1 l1) (w b2 l2)')
+    # [B, L2, B, L1, W]
+    labels_per_text = F.one_hot(
+        torch.ones(batch, text_len, batch, img_len, dtype=torch.long, device=text_x.device) * get_rank(),
+        num_classes=get_world_size()).to(text_x.dtype)
+    labels_per_text *= rearrange(pos_labels_batch_text, 'b l2 l1 -> b l2 1 l1 1') * repeat(
+        torch.eye(batch, dtype=text_x.dtype, device=image_x.device), 'b2 b1 -> b2 1 b1 1 1')
+    # [BxL2, WxBxL1]
+    labels_per_text = rearrange(labels_per_text, 'b2 l2 b1 l1 w -> (b2 l2) (w b1 l1)')
+
+    logit_scale = temperature.exp().clamp(max=100)
+
+    loss_img = soft_cross_entropy(logit_scale * logits_per_img, labels_per_img)
+    loss_text = soft_cross_entropy(logit_scale * logits_per_text, labels_per_text)
+
+    loss = 0.5 * (loss_img + loss_text)
+    return loss
+
+def vl_contrastive_loss(image_feat, text_feat, temperature=1):
+    # if image_id or text_id is None, it should be None across all GPUs
+    # image_feat = F.normalize(image_feat, dim=1)
+    # text_feat = F.normalize(text_feat, dim=1)
+    # handle normalization outside
+
+    # add the following 4 lines
+    image_feat = all_gather_grad(image_feat)
+    text_feat = all_gather_grad(text_feat)
+    
+    logits = torch.matmul(image_feat, text_feat.t())
+    logit_scale = temperature.exp().clamp(max=100)
+
+    gt = torch.arange(logits.shape[0], device=logits.device)
+    loss1 = F.cross_entropy(logit_scale * logits, gt)
+    loss2 = F.cross_entropy(logit_scale * logits.t(), gt)
+    return (loss1 + loss2) / 2 # scale it up by the number of GPUs
+
+
+def all_gather_pickle(data, device):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    # serialized to a Tensor
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to(device)
+
+    # obtain Tensor size of each rank
+    local_size = torch.LongTensor([tensor.numel()]).cuda()
+    size_list = [torch.LongTensor([0]).cuda() for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.ByteTensor(size=(max_size,)).cuda())
+    if local_size != max_size:
+        padding = torch.ByteTensor(size=(max_size - local_size,)).cuda()
+        tensor = torch.cat((tensor, padding), dim=0)
+    dist.all_gather(tensor_list, tensor)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list
+
+def all_gather_arbitary_tensor(tensor):
+    if get_world_size() > 1:
+        device = tensor.device
+        tensor_batch = all_gather_pickle(tensor.cpu(), device)
+        tensor_batch = [x.to(device) for x in tensor_batch]
+        tensor_batch[torch.distributed.get_rank()] = tensor
+        tensor_batch = torch.cat(tensor_batch, dim=0)
+    else:
+        tensor_batch = tensor
+    return tensor_batch
+
+def ql_contrastive_loss(image_feat, text_feat, temperature=1):
+    # add the following 4 lines
+    image_feat = all_gather_arbitary_tensor(image_feat)
+    text_feat = all_gather_arbitary_tensor(text_feat)
+
+    logits = torch.matmul(image_feat, text_feat.t())
+    logit_scale = temperature.exp().clamp(max=100)
+
+    gt = torch.arange(logits.shape[0], device=logits.device)
+    loss1 = F.cross_entropy(logit_scale * logits, gt)
+    loss2 = F.cross_entropy(logit_scale * logits.t(), gt)
+    return (loss1 + loss2) / 2 # scale it up by the number of GPUs
+
+def vl_similarity(image_feat, text_feat, temperature=1):
+    # Only support single GPU for now.
+    logits = torch.matmul(image_feat, text_feat.t())
+    logits = temperature.exp().clamp(max=100) * logits
+    return logits
+
+def ql_multi_contrastive_loss(image_feat, text_feat, text_hash, temperature=1):
+    # add the following 4 lines
+    image_feat = all_gather_arbitary_tensor(image_feat)
+    text_feat = all_gather_arbitary_tensor(text_feat)
+
+    text_hash_batch = all_gather_pickle(text_hash, text_feat.device)
+    text_hash_all = torch.cat(text_hash_batch)
+    
+    text_hash_all_unique = torch.unique(text_hash_all).tolist()
+    gt = torch.zeros((image_feat.shape[0], len(text_hash_all_unique)), device=text_feat.device)
+    text_hash_all = text_hash_all.tolist()
+    text_feat_unique = torch.stack([text_feat[text_hash_all.index(txt)] for txt in text_hash_all_unique])
+
+    for idx, txt in enumerate(text_hash_all):
+        gt[idx][text_hash_all_unique.index(txt)] = 1
+    
+    logits = torch.matmul(image_feat, text_feat_unique.t())
+    logits = logits*temperature.exp().clamp(max=100)
+    
+    loss_img = soft_cross_entropy(logits, gt)
+    loss_text = soft_cross_entropy(logits.t(), gt.t() / gt.t().sum(-1, keepdim=True))
+
+    loss = 0.7 * loss_img + 0.3 * loss_text
+    return loss
+
+def image_text_contrastive_loss_queue(image_feat_inp, text_feat_inp, lang_enc, training):
+    # add the following 4 lines
+    image_feat = all_gather_grad(image_feat_inp.contiguous())
+    text_feat = all_gather_grad(text_feat_inp.contiguous())
+
+    image_feat = image_feat / (image_feat.norm(dim=-1, keepdim=True) + 1e-7)
+    text_feat = text_feat / (text_feat.norm(dim=-1, keepdim=True) + 1e-7)
+
+    temperature = lang_enc.logit_scale
+    logits = torch.matmul(image_feat, text_feat.t())
+    logit_scale = temperature.exp().clamp(max=100)
+
+    gt = torch.arange(logits.shape[0], device=logits.device)
+    loss1 = F.cross_entropy(logit_scale * logits, gt)
+    loss2 = F.cross_entropy(logit_scale * logits.t(), gt)
+
+    return (loss1 + loss2) / 2 # scale it up by the number of GPUs

xdecoder/language/misc.py

@@ -0,0 +1,64 @@
+import random
+
+import nltk
+nltk.data.path.append('/mnt/data/nltk_data')
+import numpy as np
+
+from utils.constants import IMAGENET_DEFAULT_TEMPLATES
+
+
+def get_tag(tokenized, tags):
+    if not isinstance(tags, (list, tuple)):
+        tags = [tags]
+    ret = []
+    for (word, pos) in nltk.pos_tag(tokenized):
+        for tag in tags:
+            if pos == tag:
+                ret.append(word)
+    return ret
+
+def get_noun_phrase(tokenized):
+    # Taken from Su Nam Kim Paper...
+    grammar = r"""
+        NBAR:
+            {<NN.*|JJ>*<NN.*>}  # Nouns and Adjectives, terminated with Nouns
+
+        NP:
+            {<NBAR>}
+            {<NBAR><IN><NBAR>}  # Above, connected with in/of/etc...
+    """
+    chunker = nltk.RegexpParser(grammar)
+
+    chunked = chunker.parse(nltk.pos_tag(tokenized))
+    continuous_chunk = []
+    current_chunk = []
+
+    for subtree in chunked:
+        if isinstance(subtree, nltk.Tree):
+            current_chunk.append(' '.join([token for token, pos in subtree.leaves()]))
+        elif current_chunk:
+            named_entity = ' '.join(current_chunk)
+            if named_entity not in continuous_chunk:
+                continuous_chunk.append(named_entity)
+                current_chunk = []
+        else:
+            continue
+
+    return continuous_chunk
+
+def text_noun_with_prompt_all(text, phrase_prob=0.0, append_text=True):
+    tokenized = nltk.word_tokenize(text)
+    
+    if random.random() >= phrase_prob:
+        nouns = get_tag(tokenized, ['NN', 'NNS', 'NNP'])
+    else:
+        nouns = get_noun_phrase(tokenized)
+
+
+    prompt_texts = [np.random.choice(IMAGENET_DEFAULT_TEMPLATES).format(noun) for noun in nouns]
+    
+    if append_text:
+        prompt_texts += [text]
+        nouns += [text]
+    
+    return prompt_texts, nouns

xdecoder/language/registry.py

@@ -0,0 +1,13 @@
+_model_entrypoints = {}
+
+def register_model(fn):
+    module_name_split = fn.__module__.split('.')
+    model_name = module_name_split[-1]
+    _model_entrypoints[model_name] = fn
+    return fn
+
+def model_entrypoints(model_name):
+    return _model_entrypoints[model_name]
+
+def is_model(model_name):
+    return model_name in _model_entrypoints

xdecoder/language/vlpencoder.py

@@ -0,0 +1,168 @@
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from timm.models.layers import trunc_normal_
+
+from .registry import register_model
+from ..utils import configurable
+from .LangEncoder import build_tokenizer, build_lang_encoder
+from utils.misc import prompt_engineering, get_prompt_templates
+
+
+class LanguageEncoder(nn.Module):
+
+    @configurable
+    def __init__(
+        self,
+        tokenizer,
+        tokenizer_type,
+        lang_encoder,
+        lang_projection,
+        max_token_num,
+    ):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.tokenizer_type = tokenizer_type
+        self.lang_encoder = lang_encoder
+        self.lang_proj = lang_projection
+        self.max_token_num = max_token_num
+        self.logit_scale = nn.Parameter(torch.ones([]))
+
+    @classmethod
+    def from_config(cls, cfg):
+        tokenizer = build_tokenizer(cfg['MODEL']['TEXT'])
+        tokenizer_type = cfg['MODEL']['TEXT']['TOKENIZER']
+        lang_encoder = build_lang_encoder(cfg['MODEL']['TEXT'], tokenizer, cfg['VERBOSE'])
+        max_token_num = cfg['MODEL']['TEXT']['CONTEXT_LENGTH']
+        
+        dim_lang = cfg['MODEL']['TEXT']['WIDTH']
+        dim_projection = cfg['MODEL']['DIM_PROJ']
+        lang_projection = nn.Parameter(torch.empty(dim_lang, dim_projection))
+        trunc_normal_(lang_projection, std=.02)
+        
+        return {
+            "tokenizer": tokenizer,
+            "tokenizer_type": tokenizer_type,
+            "lang_encoder": lang_encoder,
+            "lang_projection": lang_projection,
+            "max_token_num": max_token_num,
+        }
+
+    def get_text_embeddings(self, class_names, name='default', is_eval=False, add_bgd=False, prompt=True, norm=True):
+        if not is_eval:
+            if prompt:
+                # randomly sample one template
+                arbitary_concepts = [
+                    prompt_engineering(class_names[label].replace('-other','').replace('-merged','').replace('-stuff',''), topk=10000, suffix='.') \
+                    for label in range(len(class_names))
+                ]
+                if add_bgd:
+                    arbitary_concepts.append("A background in coco.")
+            else:
+                arbitary_concepts = class_names
+            
+            input_ids = []
+            attention_masks = []
+            for txt in arbitary_concepts:
+                tokens = self.tokenizer(
+                    txt, padding='max_length', truncation=True, max_length=self.max_token_num, return_tensors='pt'
+                )
+                tokens['input_ids'].squeeze_()
+                tokens['attention_mask'].squeeze_()
+
+                input_ids.append(tokens['input_ids'])
+                attention_masks.append(tokens['attention_mask'])
+
+            arbitary_tokens = torch.stack(input_ids)
+            arbitary_attention_masks = torch.stack(attention_masks)
+
+            text_emb = self.forward_language((arbitary_tokens.cuda(), arbitary_attention_masks.cuda()), norm=norm)
+            setattr(self, '{}_text_embeddings'.format(name), text_emb)
+        else:
+            with torch.no_grad():
+                def extract_mean_emb(txts):
+                    tokens = self.tokenizer(
+                        txts, padding='max_length', truncation=True, max_length=self.max_token_num, return_tensors='pt'
+                    )
+                    clss_embedding = self.forward_language((tokens['input_ids'].cuda(), tokens['attention_mask'].cuda()), norm=norm)
+                    clss_embedding = clss_embedding.mean(dim=0)
+                    clss_embedding /= clss_embedding.norm()
+                    return clss_embedding
+
+                templates = get_prompt_templates()
+                clss_embeddings = []
+                if prompt:
+                    for clss in class_names:
+                        txts = [template.format(clss.replace('-other','').replace('-merged','').replace('-stuff','')) for template in templates]
+                        clss_embeddings.append(extract_mean_emb(txts))
+                else:
+                    clss_embeddings.append(extract_mean_emb(class_names))
+
+                if add_bgd:
+                    txts = ["A background in coco."]
+                    clss_embeddings.append(extract_mean_emb(txts))
+
+                text_emb = torch.stack(clss_embeddings, dim=0)
+                setattr(self, '{}_text_embeddings'.format(name), text_emb)
+
+    def get_text_token_embeddings(self, txts, name='default', token=False, norm=False):
+        if not token:
+            tokens = self.tokenizer(
+                txts, padding='max_length', truncation=True, max_length=self.max_token_num, return_tensors='pt'
+            )
+            tokens = {key: value.cuda() for key, value in tokens.items()}
+        else:
+            tokens = txts
+        token_emb, class_emb = self.forward_language_token((tokens['input_ids'], tokens['attention_mask']), norm=norm)
+        ret = {"tokens": tokens,
+                "token_emb": token_emb,
+                "class_emb": class_emb,}
+        setattr(self, '{}_token_embeddings'.format(name), ret)
+        return ret
+
+    def forward_language(self, texts, norm=True):
+        x = self.lang_encoder(*texts)
+        x = x['last_hidden_state']
+
+        if self.tokenizer_type == 'clip':
+            x = x[torch.arange(x.size(0)), texts[0].argmax(dim=-1)]
+        else:
+            x = x[:, 0]
+
+        x = x @ self.lang_proj
+        if norm:
+            x = x / (x.norm(dim=-1, keepdim=True) + 1e-7)
+        return x
+    
+    def forward_language_token(self, texts, norm=False):
+        x = self.lang_encoder(*texts)
+        token_x = x['last_hidden_state']
+
+        if self.tokenizer_type == 'clip':
+            class_x = token_x[torch.arange(token_x.size(0)), texts[0].argmax(dim=-1)]
+        else:
+            class_x = token_x[:, 0]
+
+        class_x = class_x @ self.lang_proj
+        token_x = token_x @ self.lang_proj
+
+        if norm:
+            class_x = class_x / (class_x.norm(dim=-1, keepdim=True) + 1e-7)
+            token_x = token_x / (token_x.norm(dim=-1, keepdim=True) + 1e-7)
+
+        return token_x, class_x
+    
+    def compute_similarity(self, v_emb, name='default', fake=False):
+        if fake:
+            return None
+        v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)
+        t_emb = getattr(self, '{}_text_embeddings'.format(name))
+        output = self.logit_scale.exp() * v_emb @ t_emb.unsqueeze(0).transpose(1, 2)
+        return output
+
+
+@register_model
+def get_language_model(cfg, **kwargs):
+    return LanguageEncoder(cfg)