Delete xdecoder

xdecoder/BaseModel.py

@@ -1,37 +0,0 @@
-# --------------------------------------------------------
-# X-Decoder -- Generalized Decoding for Pixel, Image, and Language
-# Copyright (c) 2022 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Xueyan Zou (xueyan@cs.wisc.edu)
-# --------------------------------------------------------
-
-import os
-import logging
-
-import torch
-import torch.nn as nn
-
-from utils.model_loading import align_and_update_state_dicts
-
-logger = logging.getLogger(__name__)
-
-
-class BaseModel(nn.Module):
-    def __init__(self, opt, module: nn.Module):
-        super(BaseModel, self).__init__()
-        self.opt = opt
-        self.model = module
-
-    def forward(self, *inputs, **kwargs):
-        outputs = self.model(*inputs, **kwargs)
-        return outputs
-
-    def save_pretrained(self, save_dir):
-        save_path = os.path.join(save_dir, 'model_state_dict.pt')
-        torch.save(self.model.state_dict(), save_path)
-
-    def from_pretrained(self, load_path):
-        state_dict = torch.load(load_path, map_location=self.opt['device'])
-        state_dict = align_and_update_state_dicts(self.model.state_dict(), state_dict)
-        self.model.load_state_dict(state_dict, strict=False)
-        return self

xdecoder/__init__.py

@@ -1,5 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from .architectures import build_model

xdecoder/architectures/__init__.py

@@ -1,2 +0,0 @@
-from .xdecoder_model import *
-from .build import build_model

xdecoder/architectures/build.py

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

xdecoder/architectures/registry.py

@@ -1,13 +0,0 @@
-_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/architectures/xdecoder_model.py

@@ -1,622 +0,0 @@
-# --------------------------------------------------------
-# X-Decoder -- Generalized Decoding for Pixel, Image, and Language
-# Copyright (c) 2022 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Xueyan Zou (xueyan@cs.wisc.edu)
-# --------------------------------------------------------
-
-import random
-from typing import Tuple
-from unicodedata import name
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-import numpy as np
-
-from .registry import register_model
-from ..utils import configurable
-from ..backbone import build_backbone, Backbone
-from ..body import build_xdecoder_head
-from ..modules import sem_seg_postprocess, bbox_postprocess
-from ..language import build_language_encoder
-from ..language.loss import vl_similarity
-
-from timm.models.layers import trunc_normal_
-from nltk.stem.lancaster import LancasterStemmer
-from detectron2.structures import Boxes, ImageList, Instances, BitMasks, BoxMode
-from detectron2.utils.memory import retry_if_cuda_oom
-from detectron2.data import MetadataCatalog
-from utils.misc import prompt_engineering
-
-st = LancasterStemmer()
-
-
-class X_Decoder_Model(nn.Module):
-    @configurable
-    def __init__(
-        self,
-        *,
-        backbone: Backbone,
-        sem_seg_head: nn.Module,
-        criterion: nn.Module,
-        losses: dict,
-        num_queries: int,
-        object_mask_threshold: float,
-        overlap_threshold: float,
-        metadata,
-        task_switch: dict,
-        phrase_prob: float,
-        size_divisibility: int,
-        sem_seg_postprocess_before_inference: bool,
-        pixel_mean: Tuple[float],
-        pixel_std: Tuple[float],
-        # inference
-        semantic_on: bool,
-        panoptic_on: bool,
-        instance_on: bool,
-        test_topk_per_image: int,
-        train_dataset_name: str,
-        retrieval_emsemble: bool,
-        backbone_dim: int,
-        dim_proj: int,
-    ):
-        super().__init__()
-        self.backbone = backbone
-        self.sem_seg_head = sem_seg_head
-        self.criterion = criterion
-        self.losses = losses
-        self.num_queries = num_queries
-        self.overlap_threshold = overlap_threshold
-        self.object_mask_threshold = object_mask_threshold
-        self.metadata = metadata
-        if size_divisibility < 0:
-            # use backbone size_divisibility if not set
-            size_divisibility = self.backbone.size_divisibility
-        self.size_divisibility = size_divisibility
-        self.sem_seg_postprocess_before_inference = sem_seg_postprocess_before_inference
-        self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False)
-        self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False)
-
-        # additional args
-        self.semantic_on = semantic_on
-        self.instance_on = instance_on
-        self.panoptic_on = panoptic_on
-
-        # caption argument
-        self.task_switch = task_switch
-        self.phrase_prob = phrase_prob
-
-        self.test_topk_per_image = test_topk_per_image
-        self.train_class_names = None
-
-        self.retrieval_emsemble = retrieval_emsemble
-        # backbone itc loss
-        if task_switch['retrieval'] and retrieval_emsemble:
-            self.backbone_proj = nn.Parameter(torch.empty(backbone_dim, dim_proj))
-            trunc_normal_(self.backbone_proj, std=.02)
-
-        if not self.semantic_on:
-            assert self.sem_seg_postprocess_before_inference
-
-    @classmethod
-    def from_config(cls, cfg):
-        enc_cfg = cfg['MODEL']['ENCODER']
-        dec_cfg = cfg['MODEL']['DECODER']
-
-        task_switch = {'bbox': dec_cfg.get('DETECTION', False),
-                       'mask': dec_cfg.get('MASK', True),
-                       'caption': dec_cfg['CAPTION'].get('ENABLED', False),
-                       'captioning': dec_cfg['CAPTIONING'].get('ENABLED', False),
-                       'retrieval': dec_cfg['RETRIEVAL'].get('ENABLED', False),
-                       'grounding': dec_cfg['GROUNDING'].get('ENABLED', False)}
-
-        # build model
-        extra = {'task_switch': task_switch}
-        backbone = build_backbone(cfg)
-        lang_encoder = build_language_encoder(cfg)        
-        sem_seg_head = build_xdecoder_head(cfg, backbone.output_shape(), lang_encoder, extra)
-
-        # Training Settings.
-        loss_weights = {}
-        matcher = None
-        losses = {}
-        weight_dict = {}        
-        grd_weight = {}
-        top_x_layers = {}
-        criterion = None
-        train_dataset_name = None
-        phrase_prob = None
-        # Loss parameters:
-        deep_supervision = None
-        no_object_weight = None
-
-        return {
-            "backbone": backbone,
-            "sem_seg_head": sem_seg_head,
-            "criterion": criterion,
-            "losses": losses,
-            "num_queries": dec_cfg['NUM_OBJECT_QUERIES'],
-            "object_mask_threshold": dec_cfg['TEST']['OBJECT_MASK_THRESHOLD'],
-            "overlap_threshold": dec_cfg['TEST']['OVERLAP_THRESHOLD'],
-            "metadata": None,
-            "size_divisibility": dec_cfg['SIZE_DIVISIBILITY'],
-            "sem_seg_postprocess_before_inference": (
-                dec_cfg['TEST']['SEM_SEG_POSTPROCESSING_BEFORE_INFERENCE']
-                or dec_cfg['TEST']['PANOPTIC_ON']
-                or dec_cfg['TEST']['INSTANCE_ON']
-            ),
-            "pixel_mean": cfg['INPUT']['PIXEL_MEAN'],
-            "pixel_std": cfg['INPUT']['PIXEL_STD'],
-            "task_switch": task_switch,
-            "phrase_prob": phrase_prob,
-            # inference
-            "semantic_on": dec_cfg['TEST']['SEMANTIC_ON'],
-            "instance_on": dec_cfg['TEST']['INSTANCE_ON'],
-            "panoptic_on": dec_cfg['TEST']['PANOPTIC_ON'],
-            "test_topk_per_image": cfg['MODEL']['DECODER']['TEST']['DETECTIONS_PER_IMAGE'],
-            "train_dataset_name": train_dataset_name,
-            "retrieval_emsemble": dec_cfg['RETRIEVAL']['ENSEMBLE'],
-            "backbone_dim": cfg['MODEL']['BACKBONE_DIM'],
-            "dim_proj": cfg['MODEL']['DIM_PROJ'],
-        }
-
-    @property
-    def device(self):
-        return self.pixel_mean.device
-
-    def forward(self, batched_inputs, mode=None):
-        if self.training:
-            assert False, "Not support trianing mode."
-        else:
-            if mode == 'retrieval':
-                return self.evaluate_retrieval(batched_inputs)
-            elif mode == 'captioning':
-                return self.evaluate_captioning(batched_inputs)
-            elif mode == 'classification':
-                return self.evaluate_classification(batched_inputs)
-            elif mode in ['grounding_phrasecut', 'grounding_refcoco']:
-                return self.evaluate_grounding(batched_inputs, mode)
-            else:
-                return self.evaluate(batched_inputs)
-        
-    def evaluate(self, batched_inputs):
-        images = [x["image"].to(self.device) for x in batched_inputs]
-        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
-        
-        images = ImageList.from_tensors(images, self.size_divisibility)
-        img_bs = images.tensor.shape[0]
-
-        targets = targets_grounding = queries_grounding = None
-        features = self.backbone(images.tensor)
-        outputs = self.sem_seg_head(features, target_queries=queries_grounding)
-
-        mask_cls_results = outputs["pred_logits"]
-        mask_pred_results = outputs["pred_masks"]
-        box_pred_results = outputs["pred_boxes"] if self.task_switch['bbox'] else [None for i in range(len(mask_pred_results))]
-        caption_pred_results = outputs["pred_captions"] if self.task_switch['caption'] else [None for i in range(len(mask_pred_results))]
-
-        # upsample masks
-        mask_pred_results = F.interpolate(
-            mask_pred_results,
-            size=(images.tensor.shape[-2], images.tensor.shape[-1]),
-            mode="bilinear",
-            align_corners=False,
-        )
-
-        input_size = mask_pred_results.shape[-2:]
-        keep_sem_bgd = self.metadata.keep_sem_bgd if hasattr(self.metadata, 'keep_sem_bgd') else False
-        del outputs
-
-        processed_results = []
-        for mask_cls_result, mask_pred_result, box_pred_result, caption_pred_result, input_per_image, image_size in zip(
-            mask_cls_results, mask_pred_results, box_pred_results, caption_pred_results, batched_inputs, images.image_sizes
-        ):
-            height = input_per_image.get("height", image_size[0])
-            width = input_per_image.get("width", image_size[1])
-            processed_results.append({})
-
-            if self.sem_seg_postprocess_before_inference:
-                mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
-                    mask_pred_result, image_size, height, width
-                )
-                mask_cls_result = mask_cls_result.to(mask_pred_result)
-
-            # semantic segmentation inference
-            if self.semantic_on:
-                r = retry_if_cuda_oom(self.semantic_inference)(mask_cls_result, mask_pred_result, keep_sem_bgd)
-                if not self.sem_seg_postprocess_before_inference:
-                    r = retry_if_cuda_oom(sem_seg_postprocess)(r, image_size, height, width)
-                processed_results[-1]["sem_seg"] = r
-
-            # panoptic segmentation inference
-            if self.panoptic_on:
-                panoptic_r = retry_if_cuda_oom(self.panoptic_inference)(mask_cls_result, mask_pred_result)
-                processed_results[-1]["panoptic_seg"] = panoptic_r
-            
-            # instance segmentation inference
-            if self.instance_on:
-                if self.task_switch['bbox']:
-                    box_pred_result = bbox_postprocess(box_pred_result, input_size, image_size, height, width)
-                instance_r = retry_if_cuda_oom(self.instance_inference)(mask_cls_result, mask_pred_result, box_pred_result)
-                processed_results[-1]["instances"] = instance_r
-            if self.task_switch['caption']:
-                processed_results[-1]["captions"] = caption_pred_result
-                processed_results[-1]["masks"] = mask_pred_result
-
-        return processed_results
-
-
-    def evaluate_retrieval(self, batched_inputs):
-        images = [x["image"].to(self.device) for x in batched_inputs]
-        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
-        images = ImageList.from_tensors(images, self.size_divisibility)
-        img_bs = images.tensor.shape[0]
-        
-        targets = targets_grounding = queries_grounding = None
-        features = self.backbone(images.tensor)
-        outputs = self.sem_seg_head(features, target_queries=queries_grounding)
-        v_emb_it = outputs['pred_captions'][:,-1]
-
-        # compute backbone score
-        if self.task_switch['retrieval'] and self.retrieval_emsemble:
-            _v_emb_it = features['res5']
-            bs,nc,_,_ = _v_emb_it.shape
-            _v_emb_it = _v_emb_it.reshape(bs,nc,-1)
-            _v_emb_it = F.adaptive_avg_pool1d(_v_emb_it, 1).reshape(bs,nc) @ self.backbone_proj
-
-        processed_results = []
-        for idx, batch_data in enumerate(batched_inputs):
-            caption_ids = []
-            t_emb_its = []
-            processed_results.append({})
-            for caption in batch_data['captions']:
-                lang_results = self.sem_seg_head.predictor.lang_encoder.get_text_token_embeddings(caption)
-                t_emb_it = lang_results['class_emb']
-                caption_ids.append(batch_data['image_id'])
-                t_emb_its.append(t_emb_it)
-
-            t_emb_it = torch.cat(t_emb_its, dim=0)
-
-            image_embeds = [v_emb_it[idx].unsqueeze(0)]
-            if self.task_switch['retrieval'] and self.retrieval_emsemble:
-                image_embeds += [_v_emb_it[idx].unsqueeze(0)]
-            caption_results = {
-                    'image_embeds': image_embeds,
-                    'text_embeds': t_emb_it,
-                    'caption_ids': caption_ids,
-                    'image_ids': batch_data['image_id'],
-                }
-            processed_results[-1]["caption"] = caption_results            
-        return processed_results
-
-    def evaluate_captioning(self, batched_inputs, extra={}):
-        images = [x["image"].to(self.device) for x in batched_inputs]
-        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
-        images = ImageList.from_tensors(images, self.size_divisibility)
-        img_bs = images.tensor.shape[0]
-
-        if not hasattr(self, 'start_token'):
-            self.start_token = torch.tensor([[49406]*77], device=self.device)
-        
-        targets = targets_grounding = queries_grounding = None
-        features = self.backbone(images.tensor)
-
-        captioning_mask = None
-        if 'captioning_mask' in batched_inputs[-1]:
-            captioning_mask = torch.cat([x['captioning_mask'] for x in batched_inputs])
-
-        extra.update({'start_token': self.start_token, 'captioning_mask': captioning_mask})
-        outputs = self.sem_seg_head(features, target_queries=queries_grounding, task='captioning_infer', extra=extra)
-
-        processed_results = []
-        for idx, batch_data in enumerate(batched_inputs):
-            processed_results.append({})
-            processed_results[-1]["captioning_token"] = outputs['pred_captionings'][idx]
-            processed_results[-1]["captioning_text"] = outputs['pred_texts'][idx].split('.')[0]
-            processed_results[-1]["image_id"] = batched_inputs[idx]['image_id']
-            
-        return processed_results
-
-    def evaluate_classification(self, batched_inputs):
-        images = [x["image"].to(self.device) for x in batched_inputs]
-        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
-        images = ImageList.from_tensors(images, self.size_divisibility)
-        img_bs = images.tensor.shape[0]
-        
-        targets = targets_grounding = queries_grounding = None
-        features = self.backbone(images.tensor)
-        outputs = self.sem_seg_head(features, target_queries=queries_grounding)
-
-        processed_results = []
-        for idx, batch_data in enumerate(batched_inputs):
-            processed_results.append({})
-            processed_results[-1]["pred_class"] = outputs['pred_logits'][idx,-1]
-        return processed_results
-
-    def evaluate_grounding_baseline(self, batched_inputs, mode):
-        images = [x["image"].to(self.device) for x in batched_inputs]
-        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
-        images = ImageList.from_tensors(images, self.size_divisibility)
-        img_bs = images.tensor.shape[0]
-        
-        targets = targets_grounding = queries_grounding = None
-        features = self.backbone(images.tensor)
-        outputs = self.sem_seg_head(features, target_queries=queries_grounding)
-
-        mask_pred_results = outputs["pred_masks"]
-        caption_pred_results = outputs["pred_captions"] if self.task_switch['caption'] else [None for i in range(len(mask_pred_results))]
-
-        # upsample masks
-        mask_pred_results = F.interpolate(
-            mask_pred_results,
-            size=(images.tensor.shape[-2], images.tensor.shape[-1]),
-            mode="bilinear",
-            align_corners=False,
-        )
-
-        processed_results = []
-        for mask_pred_result, caption_pred_result, input_per_image, image_size in zip(
-            mask_pred_results, caption_pred_results, batched_inputs, images.image_sizes
-        ):
-            height = input_per_image.get("height", image_size[0])
-            width = input_per_image.get("width", image_size[1])
-            processed_results.append({})
-
-            mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
-                mask_pred_result, image_size, height, width
-            )[:-1]
-
-            texts_all = input_per_image['groundings']['texts']
-            grd_masks = []
-            for texts in texts_all:
-                if mode == 'grounding_refcoco':
-                    self.sem_seg_head.predictor.lang_encoder.get_text_embeddings(texts, name='grounding', prompt=False, is_eval=True)
-                elif mode == 'grounding_phrasecut':
-                    self.sem_seg_head.predictor.lang_encoder.get_text_embeddings(texts, name='grounding', prompt=True, is_eval=False)
-                t_emb = getattr(self.sem_seg_head.predictor.lang_encoder, "{}_text_embeddings".format('grounding')).t()
-                v_emb = caption_pred_result[:-1]
-                v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)
-                vt_sim = v_emb @ t_emb
-                max_id = vt_sim.max(0)[1][0]
-                grd_masks += [mask_pred_result[max_id]]
-            processed_results[-1]['grounding_mask'] = torch.stack(grd_masks)
-
-        return processed_results
-
-    def evaluate_grounding(self, batched_inputs, mode):
-        images = [x["image"].to(self.device) for x in batched_inputs]
-        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
-        images = ImageList.from_tensors(images, self.size_divisibility)
-
-        extra = {}
-        # mask_pred_results = []
-        # for idx, batch_per_image in enumerate(batched_inputs):
-        #     grd_texts = batch_per_image['groundings']['texts']
-        #     grd_masks = []
-        #     for anno_text in grd_texts:
-        #         gtext = self.sem_seg_head.predictor.lang_encoder.get_text_token_embeddings([anno_text[0]], name='grounding', token=False, norm=False)
-        #         token_emb = gtext['token_emb']
-        #         tokens = gtext['tokens']
-            
-        #         grd_emb = token_emb[0][tokens['attention_mask'].bool()[0]]
-        #         extra['grounding_tokens'] = grd_emb[:,None]
-
-        #         assert len(images.tensor) == 1, "grounding evaluation only support single batch size now"
-        #         features = self.backbone(images.tensor)
-        #         outputs = self.sem_seg_head(features, extra=extra, task='grounding_eval')
-                
-        #         pred_gmasks = outputs['pred_masks'][idx,self.num_queries:2*self.num_queries-1]
-        #         v_emb = outputs['pred_captions'][idx,self.num_queries:2*self.num_queries-1]
-        #         t_emb = grd_emb[-1:]
-
-        #         t_emb = t_emb / (t_emb.norm(dim=-1, keepdim=True) + 1e-7)
-        #         v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)            
-
-        #         temperature = self.sem_seg_head.predictor.lang_encoder.logit_scale
-        #         out_prob = vl_similarity(v_emb, t_emb, temperature=temperature)
-                
-        #         matched_id = out_prob.max(0)[1]
-        #         grd_masks += [pred_gmasks[matched_id,:,:]]
-        #     mask_pred_results += [torch.cat(grd_masks)]
-
-        # comment for multi object inference.
-        mask_pred_results = []
-        for idx, batch_per_image in enumerate(batched_inputs):
-            grd_texts = batch_per_image['groundings']['texts']
-            grd_texts = [x[0] for x in grd_texts]
-
-            gtext = self.sem_seg_head.predictor.lang_encoder.get_text_token_embeddings(grd_texts, name='grounding', token=False, norm=False)
-            token_emb = gtext['token_emb']
-            tokens = gtext['tokens']
-            query_emb = token_emb[tokens['attention_mask'].bool()]
-            extra['grounding_tokens'] = query_emb[:,None]
-
-            features = self.backbone(images.tensor)
-            outputs = self.sem_seg_head(features, extra=extra, task='grounding_eval')
-
-            pred_gmasks = outputs['pred_masks'][idx,self.num_queries:2*self.num_queries-1]
-            v_emb = outputs['pred_captions'][idx,self.num_queries:2*self.num_queries-1]
-            t_emb = gtext['class_emb']
-
-            t_emb = t_emb / (t_emb.norm(dim=-1, keepdim=True) + 1e-7)
-            v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)            
-
-            temperature = self.sem_seg_head.predictor.lang_encoder.logit_scale
-            out_prob = vl_similarity(v_emb, t_emb, temperature=temperature)
-            
-            matched_id = out_prob.max(0)[1]
-            mask_pred_results += [pred_gmasks[matched_id,:,:]]
-
-        for i in range(len(mask_pred_results)):
-            # upsample masks
-            mask_pred_results[i] = F.interpolate(
-                mask_pred_results[i][None,],
-                size=(images.tensor.shape[-2], images.tensor.shape[-1]),
-                mode="bilinear",
-                align_corners=False,
-            )[0]
-
-        processed_results = []
-        for mask_pred_result, input_per_image, image_size in zip(
-            mask_pred_results, batched_inputs, images.image_sizes
-        ):
-            height = input_per_image.get("height", image_size[0])
-            width = input_per_image.get("width", image_size[1])
-            processed_results.append({})
-
-            mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
-                mask_pred_result, image_size, height, width
-            )
-            processed_results[-1]['grounding_mask'] = mask_pred_result
-
-            # compute bbox
-            # bbox = BitMasks(mask_pred_result > 0).get_bounding_boxes()
-            # bbox = BoxMode.convert(bbox.tensor, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
-            # processed_results[-1]['grounding_box'] = bbox
-
-        return processed_results
-
-    def prepare_vlp_targets(self, batched_inputs, device):
-        input_ids = []
-        attention_mask = []
-        for cnt, x in enumerate(batched_inputs):
-            captions = x['captions']
-            randid = random.randint(0, len(captions)-1)
-            input_ids += x['tokens']['input_ids'][randid:randid+1]
-            attention_mask += x['tokens']['attention_mask'][randid:randid+1]
-
-        input_ids = torch.stack(input_ids)
-        attention_mask = torch.stack(attention_mask)
-        tokens = {"input_ids": input_ids, "attention_mask": attention_mask}
-        lang_results = self.sem_seg_head.predictor.lang_encoder.get_text_token_embeddings(tokens, token=True)
-
-        target_vlp = []
-        for cnt, x in enumerate(batched_inputs):
-            target_dict = {}
-            target_dict["caption_tokens"] = lang_results['token_emb'][cnt:cnt+1]
-            target_dict["caption_proj"] = lang_results['class_emb'][cnt:cnt+1]
-            target_dict["caption_tokenids"] = lang_results['tokens']['input_ids'][cnt:cnt+1]
-            target_dict["caption_mask"] = lang_results['tokens']['attention_mask'][cnt:cnt+1]            
-            target_vlp.append(target_dict)
-        return target_vlp
-
-    def semantic_inference(self, mask_cls, mask_pred, keep_sem_bgd=False):
-        if keep_sem_bgd:
-            mask_cls = F.softmax(mask_cls, dim=-1)
-        else:
-            mask_cls = F.softmax(mask_cls, dim=-1)[..., :-1]
-        mask_pred = mask_pred.sigmoid()
-        semseg = torch.einsum("qc,qhw->chw", mask_cls, mask_pred)
-        return semseg
-
-    def panoptic_inference(self, mask_cls, mask_pred):
-        scores, labels = F.softmax(mask_cls, dim=-1).max(-1)
-        mask_pred = mask_pred.sigmoid()
-
-        keep = labels.ne(self.sem_seg_head.num_classes) & (scores > self.object_mask_threshold)
-        cur_scores = scores[keep]
-        cur_classes = labels[keep]
-        cur_masks = mask_pred[keep]
-        cur_mask_cls = mask_cls[keep]
-        cur_mask_cls = cur_mask_cls[:, :-1]
-        cur_prob_masks = cur_scores.view(-1, 1, 1) * cur_masks
-
-        h, w = cur_masks.shape[-2:]
-        panoptic_seg = torch.zeros((h, w), dtype=torch.int32, device=cur_masks.device)
-        segments_info = []
-
-        current_segment_id = 0
-
-        if cur_masks.shape[0] == 0:
-            # We didn't detect any mask :(
-            return panoptic_seg, segments_info
-        else:
-            # take argmax
-            cur_mask_ids = cur_prob_masks.argmax(0)
-            stuff_memory_list = {}
-            thing_dataset_id_to_contiguous_id = self.metadata.thing_dataset_id_to_contiguous_id if hasattr(self.metadata, 'thing_dataset_id_to_contiguous_id') else {}
-            for k in range(cur_classes.shape[0]):
-                pred_class = cur_classes[k].item()
-                isthing = pred_class in thing_dataset_id_to_contiguous_id.values()
-                mask_area = (cur_mask_ids == k).sum().item()
-                original_area = (cur_masks[k] >= 0.5).sum().item()
-                mask = (cur_mask_ids == k) & (cur_masks[k] >= 0.5)
-
-                if mask_area > 0 and original_area > 0 and mask.sum().item() > 0:
-                    if mask_area / original_area < self.overlap_threshold:
-                        continue
-
-                    # merge stuff regions
-                    if not isthing:
-                        if int(pred_class) in stuff_memory_list.keys():
-                            panoptic_seg[mask] = stuff_memory_list[int(pred_class)]
-                            continue
-                        else:
-                            stuff_memory_list[int(pred_class)] = current_segment_id + 1
-
-                    current_segment_id += 1
-                    panoptic_seg[mask] = current_segment_id
-
-                    segments_info.append(
-                        {
-                            "id": current_segment_id,
-                            "isthing": bool(isthing),
-                            "category_id": int(pred_class),
-                        }
-                    )
-            return panoptic_seg, segments_info
-
-    def instance_inference(self, mask_cls, mask_pred, box_pred):
-        # mask_pred is already processed to have the same shape as original input
-        image_size = mask_pred.shape[-2:]
-
-        # [Q, K]
-        scores = F.softmax(mask_cls, dim=-1)[:, :-1]
-        labels = torch.arange(self.sem_seg_head.num_classes, device=self.device).unsqueeze(0).repeat(self.num_queries, 1).flatten(0, 1)
-        # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
-        scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.test_topk_per_image, sorted=False)
-
-        labels_per_image = labels[topk_indices]
-        topk_indices = (topk_indices // self.sem_seg_head.num_classes)
-        # mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1)
-        mask_pred = mask_pred[topk_indices]
-        if box_pred is not None:
-            box_pred = box_pred[topk_indices]
-
-        # if this is panoptic segmentation, we only keep the "thing" classes
-        if self.panoptic_on:
-            thing_dataset_id_to_contiguous_id = self.metadata.thing_dataset_id_to_contiguous_id if hasattr(self.metadata, 'thing_dataset_id_to_contiguous_id') else {}
-            keep = torch.zeros_like(scores_per_image).bool()
-            for i, lab in enumerate(labels_per_image):
-                keep[i] = lab in thing_dataset_id_to_contiguous_id.values()
-
-            scores_per_image = scores_per_image[keep]
-            labels_per_image = labels_per_image[keep]
-            mask_pred = mask_pred[keep]
-
-            if box_pred is not None:
-                box_pred = box_pred[keep]
-
-        result = Instances(image_size)
-        # mask (before sigmoid)
-        result.pred_masks = (mask_pred > 0).float()
-        # result.pred_boxes = Boxes(torch.zeros(mask_pred.size(0), 4))
-        # Uncomment the following to get boxes from masks (this is slow)
-
-        if box_pred is not None:
-            result.pred_boxes = BitMasks(mask_pred > 0).get_bounding_boxes()
-        else:
-            result.pred_boxes = Boxes(torch.zeros(mask_pred.size(0), 4))
-
-        # calculate average mask prob
-        mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
-        result.scores = scores_per_image * mask_scores_per_image
-        result.pred_classes = labels_per_image
-
-        return result
-
-
-@register_model
-def get_segmentation_model(cfg, **kwargs):
-    return X_Decoder_Model(cfg)

xdecoder/backbone/__init__.py

@@ -1,7 +0,0 @@
-from .build import build_backbone
-
-from .resnet import *
-from .swin import *
-from .focal import *
-from .focal_dw import *
-from .backbone import *

xdecoder/backbone/backbone.py

@@ -1,51 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-import torch.nn as nn
-
-from detectron2.modeling import ShapeSpec
-
-__all__ = ["Backbone"]
-
-
-class Backbone(nn.Module):
-    """
-    Abstract base class for network backbones.
-    """
-
-    def __init__(self):
-        """
-        The `__init__` method of any subclass can specify its own set of arguments.
-        """
-        super().__init__()
-
-    def forward(self):
-        """
-        Subclasses must override this method, but adhere to the same return type.
-
-        Returns:
-            dict[str->Tensor]: mapping from feature name (e.g., "res2") to tensor
-        """
-        pass
-
-    @property
-    def size_divisibility(self) -> int:
-        """
-        Some backbones require the input height and width to be divisible by a
-        specific integer. This is typically true for encoder / decoder type networks
-        with lateral connection (e.g., FPN) for which feature maps need to match
-        dimension in the "bottom up" and "top down" paths. Set to 0 if no specific
-        input size divisibility is required.
-        """
-        return 0
-
-    def output_shape(self):
-        """
-        Returns:
-            dict[str->ShapeSpec]
-        """
-        # this is a backward-compatible default
-        return {
-            name: ShapeSpec(
-                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
-            )
-            for name in self._out_features
-        }

xdecoder/backbone/build.py

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

xdecoder/backbone/focal.py

@@ -1,692 +0,0 @@
-# --------------------------------------------------------
-# FocalNet for Semantic Segmentation
-# Copyright (c) 2022 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Jianwei Yang
-# --------------------------------------------------------
-import math
-import time
-import numpy as np
-import logging
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint as checkpoint
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
-
-from detectron2.utils.file_io import PathManager
-from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec
-
-from .registry import register_backbone
-
-logger = logging.getLogger(__name__)
-
-class Mlp(nn.Module):
-    """ Multilayer perceptron."""
-
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(hidden_features, out_features)
-        self.drop = nn.Dropout(drop)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.drop(x)
-        x = self.fc2(x)
-        x = self.drop(x)
-        return x
-
-class FocalModulation(nn.Module):
-    """ Focal Modulation
-
-    Args:
-        dim (int): Number of input channels.
-        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
-        focal_level (int): Number of focal levels
-        focal_window (int): Focal window size at focal level 1
-        focal_factor (int, default=2): Step to increase the focal window
-        use_postln (bool, default=False): Whether use post-modulation layernorm
-    """
-
-    def __init__(self, dim, proj_drop=0., focal_level=2, focal_window=7, focal_factor=2, use_postln=False, use_postln_in_modulation=False, scaling_modulator=False):
-
-        super().__init__()
-        self.dim = dim
-
-        # specific args for focalv3
-        self.focal_level = focal_level
-        self.focal_window = focal_window
-        self.focal_factor = focal_factor
-        self.use_postln_in_modulation = use_postln_in_modulation
-        self.scaling_modulator = scaling_modulator
-
-        self.f = nn.Linear(dim, 2*dim+(self.focal_level+1), bias=True)
-        self.h = nn.Conv2d(dim, dim, kernel_size=1, stride=1, padding=0, groups=1, bias=True)
-
-        self.act = nn.GELU()
-        self.proj = nn.Linear(dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-        self.focal_layers = nn.ModuleList()
-
-        if self.use_postln_in_modulation:
-            self.ln = nn.LayerNorm(dim)
-
-        for k in range(self.focal_level):
-            kernel_size = self.focal_factor*k + self.focal_window
-            self.focal_layers.append(
-                nn.Sequential(
-                    nn.Conv2d(dim, dim, kernel_size=kernel_size, stride=1, groups=dim, 
-                        padding=kernel_size//2, bias=False),
-                    nn.GELU(),
-                    )
-                )
-
-    def forward(self, x):
-        """ Forward function.
-
-        Args:
-            x: input features with shape of (B, H, W, C)
-        """
-        B, nH, nW, C = x.shape
-        x = self.f(x)
-        x = x.permute(0, 3, 1, 2).contiguous()
-        q, ctx, gates = torch.split(x, (C, C, self.focal_level+1), 1)
-        
-        ctx_all = 0
-        for l in range(self.focal_level):                     
-            ctx = self.focal_layers[l](ctx)
-            ctx_all = ctx_all + ctx*gates[:, l:l+1]
-        ctx_global = self.act(ctx.mean(2, keepdim=True).mean(3, keepdim=True))
-        ctx_all = ctx_all + ctx_global*gates[:,self.focal_level:]
-
-        if self.scaling_modulator:
-            ctx_all = ctx_all / (self.focal_level + 1)
-
-        x_out = q * self.h(ctx_all)
-        x_out = x_out.permute(0, 2, 3, 1).contiguous()
-        if self.use_postln_in_modulation:
-            x_out = self.ln(x_out)            
-        x_out = self.proj(x_out)
-        x_out = self.proj_drop(x_out)
-        return x_out
-
-class FocalModulationBlock(nn.Module):
-    """ Focal Modulation Block.
-
-    Args:
-        dim (int): Number of input channels.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
-        drop (float, optional): Dropout rate. Default: 0.0
-        drop_path (float, optional): Stochastic depth rate. Default: 0.0
-        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
-        focal_level (int): number of focal levels
-        focal_window (int): focal kernel size at level 1
-    """
-
-    def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., 
-                 act_layer=nn.GELU, norm_layer=nn.LayerNorm,
-                 focal_level=2, focal_window=9, 
-                 use_postln=False, use_postln_in_modulation=False,
-                 scaling_modulator=False, 
-                 use_layerscale=False, 
-                 layerscale_value=1e-4):
-        super().__init__()
-        self.dim = dim
-        self.mlp_ratio = mlp_ratio
-        self.focal_window = focal_window
-        self.focal_level = focal_level
-        self.use_postln = use_postln
-        self.use_layerscale = use_layerscale
-
-        self.norm1 = norm_layer(dim)
-        self.modulation = FocalModulation(
-            dim, focal_window=self.focal_window, focal_level=self.focal_level, proj_drop=drop, use_postln_in_modulation=use_postln_in_modulation, scaling_modulator=scaling_modulator
-        )            
-
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
-
-        self.H = None
-        self.W = None
-
-        self.gamma_1 = 1.0
-        self.gamma_2 = 1.0
-        if self.use_layerscale:
-            self.gamma_1 = nn.Parameter(layerscale_value * torch.ones((dim)), requires_grad=True)
-            self.gamma_2 = nn.Parameter(layerscale_value * torch.ones((dim)), requires_grad=True)
-
-    def forward(self, x):
-        """ Forward function.
-
-        Args:
-            x: Input feature, tensor size (B, H*W, C).
-            H, W: Spatial resolution of the input feature.
-        """
-        B, L, C = x.shape
-        H, W = self.H, self.W
-        assert L == H * W, "input feature has wrong size"
-
-        shortcut = x
-        if not self.use_postln:
-            x = self.norm1(x)
-        x = x.view(B, H, W, C)
-        
-        # FM
-        x = self.modulation(x).view(B, H * W, C)
-        if self.use_postln:
-            x = self.norm1(x)
-
-        # FFN
-        x = shortcut + self.drop_path(self.gamma_1 * x)
-
-        if self.use_postln:
-            x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
-        else:
-            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
-
-        return x
-
-class BasicLayer(nn.Module):
-    """ A basic focal modulation layer for one stage.
-
-    Args:
-        dim (int): Number of feature channels
-        depth (int): Depths of this stage.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
-        drop (float, optional): Dropout rate. Default: 0.0
-        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
-        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
-        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
-        focal_level (int): Number of focal levels
-        focal_window (int): Focal window size at focal level 1
-        use_conv_embed (bool): Use overlapped convolution for patch embedding or now. Default: False
-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
-    """
-
-    def __init__(self,
-                 dim,
-                 depth,
-                 mlp_ratio=4.,
-                 drop=0.,
-                 drop_path=0.,
-                 norm_layer=nn.LayerNorm,
-                 downsample=None,
-                 focal_window=9, 
-                 focal_level=2, 
-                 use_conv_embed=False,     
-                 use_postln=False,          
-                 use_postln_in_modulation=False, 
-                 scaling_modulator=False,
-                 use_layerscale=False,                   
-                 use_checkpoint=False
-        ):
-        super().__init__()
-        self.depth = depth
-        self.use_checkpoint = use_checkpoint
-
-        # build blocks
-        self.blocks = nn.ModuleList([
-            FocalModulationBlock(
-                dim=dim,
-                mlp_ratio=mlp_ratio,
-                drop=drop,
-                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
-                focal_window=focal_window, 
-                focal_level=focal_level, 
-                use_postln=use_postln, 
-                use_postln_in_modulation=use_postln_in_modulation, 
-                scaling_modulator=scaling_modulator,
-                use_layerscale=use_layerscale, 
-                norm_layer=norm_layer)
-            for i in range(depth)])
-
-        # patch merging layer
-        if downsample is not None:
-            self.downsample = downsample(
-                patch_size=2,
-                in_chans=dim, embed_dim=2*dim, 
-                use_conv_embed=use_conv_embed, 
-                norm_layer=norm_layer, 
-                is_stem=False
-            )
-
-        else:
-            self.downsample = None
-
-    def forward(self, x, H, W):
-        """ Forward function.
-
-        Args:
-            x: Input feature, tensor size (B, H*W, C).
-            H, W: Spatial resolution of the input feature.
-        """
-        for blk in self.blocks:
-            blk.H, blk.W = H, W
-            if self.use_checkpoint:
-                x = checkpoint.checkpoint(blk, x)
-            else:
-                x = blk(x)
-        if self.downsample is not None:
-            x_reshaped = x.transpose(1, 2).view(x.shape[0], x.shape[-1], H, W)
-            x_down = self.downsample(x_reshaped)   
-            x_down = x_down.flatten(2).transpose(1, 2)            
-            Wh, Ww = (H + 1) // 2, (W + 1) // 2
-            return x, H, W, x_down, Wh, Ww
-        else:
-            return x, H, W, x, H, W
-
-
-class PatchEmbed(nn.Module):
-    """ Image to Patch Embedding
-
-    Args:
-        patch_size (int): Patch token size. Default: 4.
-        in_chans (int): Number of input image channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        norm_layer (nn.Module, optional): Normalization layer. Default: None
-        use_conv_embed (bool): Whether use overlapped convolution for patch embedding. Default: False
-        is_stem (bool): Is the stem block or not. 
-    """
-
-    def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None, use_conv_embed=False, is_stem=False):
-        super().__init__()
-        patch_size = to_2tuple(patch_size)
-        self.patch_size = patch_size
-
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
-
-        if use_conv_embed:
-            # if we choose to use conv embedding, then we treat the stem and non-stem differently
-            if is_stem:
-                kernel_size = 7; padding = 2; stride = 4
-            else:
-                kernel_size = 3; padding = 1; stride = 2
-            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding)                    
-        else:
-            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
-
-        if norm_layer is not None:
-            self.norm = norm_layer(embed_dim)
-        else:
-            self.norm = None
-
-    def forward(self, x):
-        """Forward function."""
-        _, _, H, W = x.size()
-        if W % self.patch_size[1] != 0:
-            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
-        if H % self.patch_size[0] != 0:
-            x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
-
-        x = self.proj(x)  # B C Wh Ww
-        if self.norm is not None:
-            Wh, Ww = x.size(2), x.size(3)
-            x = x.flatten(2).transpose(1, 2)
-            x = self.norm(x)
-            x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
-
-        return x
-
-
-class FocalNet(nn.Module):
-    """ FocalNet backbone.
-
-    Args:
-        pretrain_img_size (int): Input image size for training the pretrained model,
-            used in absolute postion embedding. Default 224.
-        patch_size (int | tuple(int)): Patch size. Default: 4.
-        in_chans (int): Number of input image channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        depths (tuple[int]): Depths of each Swin Transformer stage.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
-        drop_rate (float): Dropout rate.
-        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
-        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
-        patch_norm (bool): If True, add normalization after patch embedding. Default: True.
-        out_indices (Sequence[int]): Output from which stages.
-        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
-            -1 means not freezing any parameters.
-        focal_levels (Sequence[int]): Number of focal levels at four stages
-        focal_windows (Sequence[int]): Focal window sizes at first focal level at four stages
-        use_conv_embed (bool): Whether use overlapped convolution for patch embedding
-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
-    """
-
-    def __init__(self,
-                 pretrain_img_size=1600,
-                 patch_size=4,
-                 in_chans=3,
-                 embed_dim=96,
-                 depths=[2, 2, 6, 2],
-                 mlp_ratio=4.,
-                 drop_rate=0.,
-                 drop_path_rate=0.2,
-                 norm_layer=nn.LayerNorm,
-                 patch_norm=True,
-                 out_indices=[0, 1, 2, 3],
-                 frozen_stages=-1,
-                 focal_levels=[2,2,2,2], 
-                 focal_windows=[9,9,9,9],
-                 use_conv_embed=False, 
-                 use_postln=False, 
-                 use_postln_in_modulation=False, 
-                 scaling_modulator=False,
-                 use_layerscale=False, 
-                 use_checkpoint=False, 
-        ):
-        super().__init__()
-
-        self.pretrain_img_size = pretrain_img_size
-        self.num_layers = len(depths)
-        self.embed_dim = embed_dim
-        self.patch_norm = patch_norm
-        self.out_indices = out_indices
-        self.frozen_stages = frozen_stages
-
-        # split image into non-overlapping patches
-        self.patch_embed = PatchEmbed(
-            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
-            norm_layer=norm_layer if self.patch_norm else None, 
-            use_conv_embed=use_conv_embed, is_stem=True)
-
-        self.pos_drop = nn.Dropout(p=drop_rate)
-
-        # stochastic depth
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
-
-        # build layers
-        self.layers = nn.ModuleList()
-        for i_layer in range(self.num_layers):
-            layer = BasicLayer(
-                dim=int(embed_dim * 2 ** i_layer),
-                depth=depths[i_layer],
-                mlp_ratio=mlp_ratio,
-                drop=drop_rate,
-                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
-                norm_layer=norm_layer,
-                downsample=PatchEmbed if (i_layer < self.num_layers - 1) else None,
-                focal_window=focal_windows[i_layer], 
-                focal_level=focal_levels[i_layer], 
-                use_conv_embed=use_conv_embed,
-                use_postln=use_postln, 
-                use_postln_in_modulation=use_postln_in_modulation,
-                scaling_modulator=scaling_modulator,
-                use_layerscale=use_layerscale, 
-                use_checkpoint=use_checkpoint)
-            self.layers.append(layer)
-
-        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
-        self.num_features = num_features
-
-        # add a norm layer for each output
-        for i_layer in out_indices:
-            layer = norm_layer(num_features[i_layer])
-            layer_name = f'norm{i_layer}'
-            self.add_module(layer_name, layer)
-
-        self._freeze_stages()
-
-    def _freeze_stages(self):
-        if self.frozen_stages >= 0:
-            self.patch_embed.eval()
-            for param in self.patch_embed.parameters():
-                param.requires_grad = False
-
-        if self.frozen_stages >= 2:
-            self.pos_drop.eval()
-            for i in range(0, self.frozen_stages - 1):
-                m = self.layers[i]
-                m.eval()
-                for param in m.parameters():
-                    param.requires_grad = False
-
-    def init_weights(self, pretrained=None):
-        """Initialize the weights in backbone.
-
-        Args:
-            pretrained (str, optional): Path to pre-trained weights.
-                Defaults to None.
-        """
-
-        def _init_weights(m):
-            if isinstance(m, nn.Linear):
-                trunc_normal_(m.weight, std=.02)
-                if isinstance(m, nn.Linear) and m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.LayerNorm):
-                nn.init.constant_(m.bias, 0)
-                nn.init.constant_(m.weight, 1.0)
-
-        if isinstance(pretrained, str):
-            self.apply(_init_weights)
-            logger = get_root_logger()
-            load_checkpoint(self, pretrained, strict=False, logger=logger)
-        elif pretrained is None:
-            self.apply(_init_weights)
-        else:
-            raise TypeError('pretrained must be a str or None')
-
-    def load_weights(self, pretrained_dict=None, pretrained_layers=[], verbose=True):
-        model_dict = self.state_dict()
-
-        missed_dict = [k for k in model_dict.keys() if k not in pretrained_dict]
-        logger.info(f'=> Missed keys {missed_dict}')
-        unexpected_dict = [k for k in pretrained_dict.keys() if k not in model_dict]
-        logger.info(f'=> Unexpected keys {unexpected_dict}')
-
-        pretrained_dict = {
-            k: v for k, v in pretrained_dict.items()
-            if k in model_dict.keys()
-        }
-        
-        need_init_state_dict = {}
-        for k, v in pretrained_dict.items():
-            need_init = (
-                (
-                    k.split('.')[0] in pretrained_layers
-                    or pretrained_layers[0] == '*'
-                )
-                and 'relative_position_index' not in k
-                and 'attn_mask' not in k
-            )
-
-            if need_init:
-                # if verbose:
-                #     logger.info(f'=> init {k} from {pretrained}')
-
-                if ('pool_layers' in k) or ('focal_layers' in k) and v.size() != model_dict[k].size():
-                    table_pretrained = v
-                    table_current = model_dict[k]
-                    fsize1 = table_pretrained.shape[2]
-                    fsize2 = table_current.shape[2]
-
-                    # NOTE: different from interpolation used in self-attention, we use padding or clipping for focal conv
-                    if fsize1 < fsize2:
-                        table_pretrained_resized = torch.zeros(table_current.shape)
-                        table_pretrained_resized[:, :, (fsize2-fsize1)//2:-(fsize2-fsize1)//2, (fsize2-fsize1)//2:-(fsize2-fsize1)//2] = table_pretrained
-                        v = table_pretrained_resized
-                    elif fsize1 > fsize2:
-                        table_pretrained_resized = table_pretrained[:, :, (fsize1-fsize2)//2:-(fsize1-fsize2)//2, (fsize1-fsize2)//2:-(fsize1-fsize2)//2]
-                        v = table_pretrained_resized
-
-
-                if ("modulation.f" in k or "pre_conv" in k): 
-                    table_pretrained = v
-                    table_current = model_dict[k]
-                    if table_pretrained.shape != table_current.shape:
-                        if len(table_pretrained.shape) == 2:
-                            dim = table_pretrained.shape[1]
-                            assert table_current.shape[1] == dim
-                            L1 = table_pretrained.shape[0]
-                            L2 = table_current.shape[0]
-
-                            if L1 < L2:
-                                table_pretrained_resized = torch.zeros(table_current.shape)
-                                # copy for linear project
-                                table_pretrained_resized[:2*dim] = table_pretrained[:2*dim]
-                                # copy for global token gating
-                                table_pretrained_resized[-1] = table_pretrained[-1]
-                                # copy for first multiple focal levels
-                                table_pretrained_resized[2*dim:2*dim+(L1-2*dim-1)] = table_pretrained[2*dim:-1]
-                                # reassign pretrained weights
-                                v = table_pretrained_resized
-                            elif L1 > L2:
-                                raise NotImplementedError
-                        elif len(table_pretrained.shape) == 1:
-                            dim = table_pretrained.shape[0]
-                            L1 = table_pretrained.shape[0]
-                            L2 = table_current.shape[0]
-                            if L1 < L2:
-                                table_pretrained_resized = torch.zeros(table_current.shape)
-                                # copy for linear project
-                                table_pretrained_resized[:dim] = table_pretrained[:dim]
-                                # copy for global token gating
-                                table_pretrained_resized[-1] = table_pretrained[-1]
-                                # copy for first multiple focal levels
-                                # table_pretrained_resized[dim:2*dim+(L1-2*dim-1)] = table_pretrained[2*dim:-1]
-                                # reassign pretrained weights
-                                v = table_pretrained_resized
-                            elif L1 > L2:
-                                raise NotImplementedError    
-
-                need_init_state_dict[k] = v
-        
-        self.load_state_dict(need_init_state_dict, strict=False)
-
-
-    def forward(self, x):
-        """Forward function."""
-        tic = time.time()
-        x = self.patch_embed(x)
-        Wh, Ww = x.size(2), x.size(3)
-
-        x = x.flatten(2).transpose(1, 2)
-        x = self.pos_drop(x)
-
-        outs = {}
-        for i in range(self.num_layers):
-            layer = self.layers[i]
-            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
-            if i in self.out_indices:
-                norm_layer = getattr(self, f'norm{i}')
-                x_out = norm_layer(x_out)
-
-                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
-                outs["res{}".format(i + 2)] = out
-                
-        if len(self.out_indices) == 0:
-            outs["res5"] = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
-
-        toc = time.time()
-        return outs
-
-    def train(self, mode=True):
-        """Convert the model into training mode while keep layers freezed."""
-        super(FocalNet, self).train(mode)
-        self._freeze_stages()
-
-
-class D2FocalNet(FocalNet, Backbone):
-    def __init__(self, cfg, input_shape):
-
-        pretrain_img_size = cfg['BACKBONE']['FOCAL']['PRETRAIN_IMG_SIZE']
-        patch_size = cfg['BACKBONE']['FOCAL']['PATCH_SIZE']
-        in_chans = 3
-        embed_dim = cfg['BACKBONE']['FOCAL']['EMBED_DIM']
-        depths = cfg['BACKBONE']['FOCAL']['DEPTHS']
-        mlp_ratio = cfg['BACKBONE']['FOCAL']['MLP_RATIO']
-        drop_rate = cfg['BACKBONE']['FOCAL']['DROP_RATE']
-        drop_path_rate = cfg['BACKBONE']['FOCAL']['DROP_PATH_RATE']
-        norm_layer = nn.LayerNorm
-        patch_norm = cfg['BACKBONE']['FOCAL']['PATCH_NORM']
-        use_checkpoint = cfg['BACKBONE']['FOCAL']['USE_CHECKPOINT']
-        out_indices = cfg['BACKBONE']['FOCAL']['OUT_INDICES']
-        scaling_modulator = cfg['BACKBONE']['FOCAL'].get('SCALING_MODULATOR', False)
-
-        super().__init__(
-            pretrain_img_size,
-            patch_size,
-            in_chans,
-            embed_dim,
-            depths,
-            mlp_ratio,
-            drop_rate,
-            drop_path_rate,
-            norm_layer,
-            patch_norm,
-            out_indices,
-            focal_levels=cfg['BACKBONE']['FOCAL']['FOCAL_LEVELS'],
-            focal_windows=cfg['BACKBONE']['FOCAL']['FOCAL_WINDOWS'],   
-            use_conv_embed=cfg['BACKBONE']['FOCAL']['USE_CONV_EMBED'],    
-            use_postln=cfg['BACKBONE']['FOCAL']['USE_POSTLN'],       
-            use_postln_in_modulation=cfg['BACKBONE']['FOCAL']['USE_POSTLN_IN_MODULATION'], 
-            scaling_modulator=scaling_modulator,
-            use_layerscale=cfg['BACKBONE']['FOCAL']['USE_LAYERSCALE'], 
-            use_checkpoint=use_checkpoint,
-        )
-
-        self._out_features = cfg['BACKBONE']['FOCAL']['OUT_FEATURES']
-
-        self._out_feature_strides = {
-            "res2": 4,
-            "res3": 8,
-            "res4": 16,
-            "res5": 32,
-        }
-        self._out_feature_channels = {
-            "res2": self.num_features[0],
-            "res3": self.num_features[1],
-            "res4": self.num_features[2],
-            "res5": self.num_features[3],
-        }
-
-    def forward(self, x):
-        """
-        Args:
-            x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
-        Returns:
-            dict[str->Tensor]: names and the corresponding features
-        """
-        assert (
-            x.dim() == 4
-        ), f"SwinTransformer takes an input of shape (N, C, H, W). Got {x.shape} instead!"
-        outputs = {}
-        y = super().forward(x)
-        for k in y.keys():
-            if k in self._out_features:
-                outputs[k] = y[k]
-        return outputs
-
-    def output_shape(self):
-        return {
-            name: ShapeSpec(
-                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
-            )
-            for name in self._out_features
-        }
-
-    @property
-    def size_divisibility(self):
-        return 32
-
-@register_backbone
-def get_focal_backbone(cfg):
-    focal = D2FocalNet(cfg['MODEL'], 224)    
-
-    if cfg['MODEL']['BACKBONE']['LOAD_PRETRAINED'] is True:
-        filename = cfg['MODEL']['BACKBONE']['PRETRAINED']
-        logger.info(f'=> init from {filename}')
-        with PathManager.open(filename, "rb") as f:
-            ckpt = torch.load(f)['model']
-        focal.load_weights(ckpt, cfg['MODEL']['BACKBONE']['FOCAL'].get('PRETRAINED_LAYERS', ['*']), cfg['VERBOSE'])
-
-    return focal

xdecoder/backbone/focal_dw.py

@@ -1,789 +0,0 @@
-# --------------------------------------------------------
-# FocalNet for Semantic Segmentation
-# Copyright (c) 2022 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Jianwei Yang
-# --------------------------------------------------------
-import math
-import time
-import numpy as np
-import logging
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint as checkpoint
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
-
-from detectron2.utils.file_io import PathManager
-from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec
-
-from .registry import register_backbone
-
-logger = logging.getLogger(__name__)
-
-class Mlp(nn.Module):
-    """ Multilayer perceptron."""
-
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(hidden_features, out_features)
-        self.drop = nn.Dropout(drop)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.drop(x)
-        x = self.fc2(x)
-        x = self.drop(x)
-        return x
-
-class FocalModulation(nn.Module):
-    """ Focal Modulation
-
-    Args:
-        dim (int): Number of input channels.
-        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
-        focal_level (int): Number of focal levels
-        focal_window (int): Focal window size at focal level 1
-        focal_factor (int, default=2): Step to increase the focal window
-        use_postln (bool, default=False): Whether use post-modulation layernorm
-    """
-
-    def __init__(self, dim, proj_drop=0., focal_level=2, focal_window=7, focal_factor=2, use_postln=False, use_postln_in_modulation=False, scaling_modulator=False):
-
-        super().__init__()
-        self.dim = dim
-
-        # specific args for focalv3
-        self.focal_level = focal_level
-        self.focal_window = focal_window
-        self.focal_factor = focal_factor
-        self.use_postln_in_modulation = use_postln_in_modulation
-        self.scaling_modulator = scaling_modulator
-
-        self.f = nn.Linear(dim, 2*dim+(self.focal_level+1), bias=True)
-        self.h = nn.Conv2d(dim, dim, kernel_size=1, stride=1, padding=0, groups=1, bias=True)
-
-        self.act = nn.GELU()
-        self.proj = nn.Linear(dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-        self.focal_layers = nn.ModuleList()
-
-        if self.use_postln_in_modulation:
-            self.ln = nn.LayerNorm(dim)
-
-        for k in range(self.focal_level):
-            kernel_size = self.focal_factor*k + self.focal_window
-            self.focal_layers.append(
-                nn.Sequential(
-                    nn.Conv2d(dim, dim, kernel_size=kernel_size, stride=1, groups=dim, 
-                        padding=kernel_size//2, bias=False),
-                    nn.GELU(),
-                    )
-                )
-
-    def forward(self, x):
-        """ Forward function.
-
-        Args:
-            x: input features with shape of (B, H, W, C)
-        """
-        B, nH, nW, C = x.shape
-        x = self.f(x)
-        x = x.permute(0, 3, 1, 2).contiguous()
-        q, ctx, gates = torch.split(x, (C, C, self.focal_level+1), 1)
-        
-        ctx_all = 0
-        for l in range(self.focal_level):                     
-            ctx = self.focal_layers[l](ctx)
-            ctx_all = ctx_all + ctx*gates[:, l:l+1]
-        ctx_global = self.act(ctx.mean(2, keepdim=True).mean(3, keepdim=True))
-        ctx_all = ctx_all + ctx_global*gates[:,self.focal_level:]
-
-        if self.scaling_modulator:
-            ctx_all = ctx_all / (self.focal_level + 1)
-
-        x_out = q * self.h(ctx_all)
-        x_out = x_out.permute(0, 2, 3, 1).contiguous()
-        if self.use_postln_in_modulation:
-            x_out = self.ln(x_out)            
-        x_out = self.proj(x_out)
-        x_out = self.proj_drop(x_out)
-        return x_out
-
-class FocalModulationBlock(nn.Module):
-    """ Focal Modulation Block.
-
-    Args:
-        dim (int): Number of input channels.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
-        drop (float, optional): Dropout rate. Default: 0.0
-        drop_path (float, optional): Stochastic depth rate. Default: 0.0
-        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
-        focal_level (int): number of focal levels
-        focal_window (int): focal kernel size at level 1
-    """
-
-    def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., 
-                 act_layer=nn.GELU, norm_layer=nn.LayerNorm,
-                 focal_level=2, focal_window=9, 
-                 use_postln=False, use_postln_in_modulation=False,
-                 scaling_modulator=False, 
-                 use_layerscale=False, 
-                 layerscale_value=1e-4):
-        super().__init__()
-        self.dim = dim
-        self.mlp_ratio = mlp_ratio
-        self.focal_window = focal_window
-        self.focal_level = focal_level
-        self.use_postln = use_postln
-        self.use_layerscale = use_layerscale
-
-        self.dw1 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim)
-        self.norm1 = norm_layer(dim)
-        self.modulation = FocalModulation(
-            dim, focal_window=self.focal_window, focal_level=self.focal_level, proj_drop=drop, use_postln_in_modulation=use_postln_in_modulation, scaling_modulator=scaling_modulator
-        )            
-
-        self.dw2 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim)
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
-
-        self.H = None
-        self.W = None
-
-        self.gamma_1 = 1.0
-        self.gamma_2 = 1.0
-        if self.use_layerscale:
-            self.gamma_1 = nn.Parameter(layerscale_value * torch.ones((dim)), requires_grad=True)
-            self.gamma_2 = nn.Parameter(layerscale_value * torch.ones((dim)), requires_grad=True)
-
-    def forward(self, x):
-        """ Forward function.
-
-        Args:
-            x: Input feature, tensor size (B, H*W, C).
-            H, W: Spatial resolution of the input feature.
-        """
-        B, L, C = x.shape
-        H, W = self.H, self.W
-        assert L == H * W, "input feature has wrong size"
-
-        x = x.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
-        x = x + self.dw1(x)
-        x = x.permute(0, 2, 3, 1).contiguous().view(B, L, C)
-
-        shortcut = x
-        if not self.use_postln:
-            x = self.norm1(x)
-        x = x.view(B, H, W, C)
-        
-        # FM
-        x = self.modulation(x).view(B, H * W, C)
-        x = shortcut + self.drop_path(self.gamma_1 * x)
-        if self.use_postln:
-            x = self.norm1(x)
-
-        x = x.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
-        x = x + self.dw2(x)
-        x = x.permute(0, 2, 3, 1).contiguous().view(B, L, C)
-
-        if not self.use_postln:
-            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))        
-        else:
-            x = x + self.drop_path(self.gamma_2 * self.mlp(x))
-            x = self.norm2(x)
-
-        return x
-
-class BasicLayer(nn.Module):
-    """ A basic focal modulation layer for one stage.
-
-    Args:
-        dim (int): Number of feature channels
-        depth (int): Depths of this stage.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
-        drop (float, optional): Dropout rate. Default: 0.0
-        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
-        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
-        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
-        focal_level (int): Number of focal levels
-        focal_window (int): Focal window size at focal level 1
-        use_conv_embed (bool): Use overlapped convolution for patch embedding or now. Default: False
-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
-    """
-
-    def __init__(self,
-                 dim,
-                 depth,
-                 mlp_ratio=4.,
-                 drop=0.,
-                 drop_path=0.,
-                 norm_layer=nn.LayerNorm,
-                 downsample=None,
-                 focal_window=9, 
-                 focal_level=2, 
-                 use_conv_embed=False,     
-                 use_postln=False,          
-                 use_postln_in_modulation=False, 
-                 scaling_modulator=False,
-                 use_layerscale=False,                   
-                 use_checkpoint=False, 
-                 use_pre_norm=False, 
-        ):
-        super().__init__()
-        self.depth = depth
-        self.use_checkpoint = use_checkpoint
-
-        # build blocks
-        self.blocks = nn.ModuleList([
-            FocalModulationBlock(
-                dim=dim,
-                mlp_ratio=mlp_ratio,
-                drop=drop,
-                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
-                focal_window=focal_window, 
-                focal_level=focal_level, 
-                use_postln=use_postln, 
-                use_postln_in_modulation=use_postln_in_modulation, 
-                scaling_modulator=scaling_modulator,
-                use_layerscale=use_layerscale, 
-                norm_layer=norm_layer)
-            for i in range(depth)])
-
-        # patch merging layer
-        if downsample is not None:
-            self.downsample = downsample(
-                patch_size=2,
-                in_chans=dim, embed_dim=2*dim, 
-                use_conv_embed=use_conv_embed, 
-                norm_layer=norm_layer, 
-                is_stem=False, 
-                use_pre_norm=use_pre_norm
-            )
-
-        else:
-            self.downsample = None
-
-    def forward(self, x, H, W):
-        """ Forward function.
-
-        Args:
-            x: Input feature, tensor size (B, H*W, C).
-            H, W: Spatial resolution of the input feature.
-        """
-        for blk in self.blocks:
-            blk.H, blk.W = H, W
-            if self.use_checkpoint:
-                x = checkpoint.checkpoint(blk, x)
-            else:
-                x = blk(x)
-        if self.downsample is not None:
-            x_reshaped = x.transpose(1, 2).view(x.shape[0], x.shape[-1], H, W)
-            x_down = self.downsample(x_reshaped)   
-            x_down = x_down.flatten(2).transpose(1, 2)            
-            Wh, Ww = (H + 1) // 2, (W + 1) // 2
-            return x, H, W, x_down, Wh, Ww
-        else:
-            return x, H, W, x, H, W
-
-
-# class PatchEmbed(nn.Module):
-#     r""" Image to Patch Embedding
-
-#     Args:
-#         img_size (int): Image size.  Default: 224.
-#         patch_size (int): Patch token size. Default: 4.
-#         in_chans (int): Number of input image channels. Default: 3.
-#         embed_dim (int): Number of linear projection output channels. Default: 96.
-#         norm_layer (nn.Module, optional): Normalization layer. Default: None
-#     """
-
-#     def __init__(self, img_size=(224, 224), patch_size=4, in_chans=3, embed_dim=96, 
-#         use_conv_embed=False, norm_layer=None, is_stem=False, use_pre_norm=False):
-#         super().__init__()
-#         patch_size = to_2tuple(patch_size)
-#         patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
-#         self.img_size = img_size
-#         self.patch_size = patch_size
-#         self.patches_resolution = patches_resolution
-#         self.num_patches = patches_resolution[0] * patches_resolution[1]
-
-#         self.in_chans = in_chans
-#         self.embed_dim = embed_dim
-#         self.use_pre_norm = use_pre_norm
-
-#         if use_conv_embed:
-#             # if we choose to use conv embedding, then we treat the stem and non-stem differently
-#             if is_stem:
-#                 kernel_size = 7; padding = 3; stride = 4
-#             else:
-#                 kernel_size = 3; padding = 1; stride = 2
-#             self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding)
-#         else:
-#             self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
-        
-#         if self.use_pre_norm:
-#             if norm_layer is not None:
-#                 self.norm = norm_layer(in_chans)
-#             else:
-#                 self.norm = None
-#         else:
-#             if norm_layer is not None:
-#                 self.norm = norm_layer(embed_dim)
-#             else:
-#                 self.norm = None
-
-#     def forward(self, x):
-#         B, C, H, W = x.shape
-#         # FIXME look at relaxing size constraints
-#         assert H == self.img_size[0] and W == self.img_size[1], \
-#             f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
-        
-#         if self.use_pre_norm:
-#             if self.norm is not None:
-#                 x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
-#                 x = self.norm(x).transpose(1, 2).view(B, C, H, W)
-#             x = self.proj(x).flatten(2).transpose(1, 2)
-#         else:
-#             x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C
-#             if self.norm is not None:
-#                 x = self.norm(x)
-#         return x
-
-#     def flops(self):
-#         Ho, Wo = self.patches_resolution
-#         flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
-#         if self.norm is not None:
-#             flops += Ho * Wo * self.embed_dim
-#         return flops
-
-class PatchEmbed(nn.Module):
-    """ Image to Patch Embedding
-
-    Args:
-        patch_size (int): Patch token size. Default: 4.
-        in_chans (int): Number of input image channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        norm_layer (nn.Module, optional): Normalization layer. Default: None
-        use_conv_embed (bool): Whether use overlapped convolution for patch embedding. Default: False
-        is_stem (bool): Is the stem block or not. 
-    """
-
-    def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None, use_conv_embed=False, is_stem=False, use_pre_norm=False):
-        super().__init__()
-        patch_size = to_2tuple(patch_size)
-        self.patch_size = patch_size
-
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
-        self.use_pre_norm = use_pre_norm
-
-        if use_conv_embed:
-            # if we choose to use conv embedding, then we treat the stem and non-stem differently
-            if is_stem:
-                kernel_size = 7; padding = 3; stride = 4
-            else:
-                kernel_size = 3; padding = 1; stride = 2
-            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding)                    
-        else:
-            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
-
-        if self.use_pre_norm:
-            if norm_layer is not None:
-                self.norm = norm_layer(in_chans)
-            else:
-                self.norm = None       
-        else:
-            if norm_layer is not None:
-                self.norm = norm_layer(embed_dim)
-            else:
-                self.norm = None
-
-    def forward(self, x):
-        """Forward function."""
-        B, C, H, W = x.size()
-        if W % self.patch_size[1] != 0:
-            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
-        if H % self.patch_size[0] != 0:
-            x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
-
-        if self.use_pre_norm:
-            if self.norm is not None:
-                x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
-                x = self.norm(x).transpose(1, 2).view(B, C, H, W)
-            x = self.proj(x)
-        else:
-            x = self.proj(x)  # B C Wh Ww
-            if self.norm is not None:
-                Wh, Ww = x.size(2), x.size(3)
-                x = x.flatten(2).transpose(1, 2)
-                x = self.norm(x)
-                x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
-
-        return x
-
-
-class FocalNet(nn.Module):
-    """ FocalNet backbone.
-
-    Args:
-        pretrain_img_size (int): Input image size for training the pretrained model,
-            used in absolute postion embedding. Default 224.
-        patch_size (int | tuple(int)): Patch size. Default: 4.
-        in_chans (int): Number of input image channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        depths (tuple[int]): Depths of each Swin Transformer stage.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
-        drop_rate (float): Dropout rate.
-        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
-        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
-        patch_norm (bool): If True, add normalization after patch embedding. Default: True.
-        out_indices (Sequence[int]): Output from which stages.
-        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
-            -1 means not freezing any parameters.
-        focal_levels (Sequence[int]): Number of focal levels at four stages
-        focal_windows (Sequence[int]): Focal window sizes at first focal level at four stages
-        use_conv_embed (bool): Whether use overlapped convolution for patch embedding
-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
-    """
-
-    def __init__(self,
-                 pretrain_img_size=1600,
-                 patch_size=4,
-                 in_chans=3,
-                 embed_dim=96,
-                 depths=[2, 2, 6, 2],
-                 mlp_ratio=4.,
-                 drop_rate=0.,
-                 drop_path_rate=0.2,
-                 norm_layer=nn.LayerNorm,
-                 patch_norm=True,
-                 out_indices=[0, 1, 2, 3],
-                 frozen_stages=-1,
-                 focal_levels=[2,2,2,2], 
-                 focal_windows=[9,9,9,9],
-                 use_pre_norms=[False, False, False, False], 
-                 use_conv_embed=False, 
-                 use_postln=False, 
-                 use_postln_in_modulation=False, 
-                 scaling_modulator=False,
-                 use_layerscale=False, 
-                 use_checkpoint=False, 
-        ):
-        super().__init__()
-
-        self.pretrain_img_size = pretrain_img_size
-        self.num_layers = len(depths)
-        self.embed_dim = embed_dim
-        self.patch_norm = patch_norm
-        self.out_indices = out_indices
-        self.frozen_stages = frozen_stages
-
-        # split image into non-overlapping patches
-        self.patch_embed = PatchEmbed(
-            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
-            norm_layer=norm_layer if self.patch_norm else None, 
-            use_conv_embed=use_conv_embed, is_stem=True, use_pre_norm=False)
-
-        self.pos_drop = nn.Dropout(p=drop_rate)
-
-        # stochastic depth
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
-
-        # build layers
-        self.layers = nn.ModuleList()
-        for i_layer in range(self.num_layers):
-            layer = BasicLayer(
-                dim=int(embed_dim * 2 ** i_layer),
-                depth=depths[i_layer],
-                mlp_ratio=mlp_ratio,
-                drop=drop_rate,
-                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
-                norm_layer=norm_layer,
-                downsample=PatchEmbed if (i_layer < self.num_layers - 1) else None,
-                focal_window=focal_windows[i_layer], 
-                focal_level=focal_levels[i_layer], 
-                use_pre_norm=use_pre_norms[i_layer], 
-                use_conv_embed=use_conv_embed,
-                use_postln=use_postln, 
-                use_postln_in_modulation=use_postln_in_modulation,
-                scaling_modulator=scaling_modulator,
-                use_layerscale=use_layerscale, 
-                use_checkpoint=use_checkpoint)
-            self.layers.append(layer)
-
-        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
-        self.num_features = num_features        
-        # self.norm = norm_layer(num_features[-1])
-
-        # add a norm layer for each output
-        for i_layer in self.out_indices:
-            layer = norm_layer(num_features[i_layer])
-            layer_name = f'norm{i_layer}'
-            self.add_module(layer_name, layer)
-
-        self._freeze_stages()
-
-    def _freeze_stages(self):
-        if self.frozen_stages >= 0:
-            self.patch_embed.eval()
-            for param in self.patch_embed.parameters():
-                param.requires_grad = False
-
-        if self.frozen_stages >= 2:
-            self.pos_drop.eval()
-            for i in range(0, self.frozen_stages - 1):
-                m = self.layers[i]
-                m.eval()
-                for param in m.parameters():
-                    param.requires_grad = False
-
-    def init_weights(self, pretrained=None):
-        """Initialize the weights in backbone.
-
-        Args:
-            pretrained (str, optional): Path to pre-trained weights.
-                Defaults to None.
-        """
-
-        def _init_weights(m):
-            if isinstance(m, nn.Linear):
-                trunc_normal_(m.weight, std=.02)
-                if isinstance(m, nn.Linear) and m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.LayerNorm):
-                nn.init.constant_(m.bias, 0)
-                nn.init.constant_(m.weight, 1.0)
-
-        if isinstance(pretrained, str):
-            self.apply(_init_weights)
-            logger = get_root_logger()
-            load_checkpoint(self, pretrained, strict=False, logger=logger)
-        elif pretrained is None:
-            self.apply(_init_weights)
-        else:
-            raise TypeError('pretrained must be a str or None')
-
-    def load_weights(self, pretrained_dict=None, pretrained_layers=[], verbose=True):
-        model_dict = self.state_dict()
-
-        missed_dict = [k for k in model_dict.keys() if k not in pretrained_dict]
-        logger.info(f'=> Missed keys {missed_dict}')
-        unexpected_dict = [k for k in pretrained_dict.keys() if k not in model_dict]
-        logger.info(f'=> Unexpected keys {unexpected_dict}')
-
-        pretrained_dict = {
-            k: v for k, v in pretrained_dict.items()
-            if k in model_dict.keys()
-        }
-        
-        need_init_state_dict = {}
-        for k, v in pretrained_dict.items():
-            need_init = (
-                (
-                    k.split('.')[0] in pretrained_layers
-                    or pretrained_layers[0] == '*'
-                )
-                and 'relative_position_index' not in k
-                and 'attn_mask' not in k
-            )
-
-            if need_init:
-                # if verbose:
-                #     logger.info(f'=> init {k} from {pretrained}')
-
-                if ('pool_layers' in k) or ('focal_layers' in k) and v.size() != model_dict[k].size():
-                    table_pretrained = v
-                    table_current = model_dict[k]
-                    fsize1 = table_pretrained.shape[2]
-                    fsize2 = table_current.shape[2]
-
-                    # NOTE: different from interpolation used in self-attention, we use padding or clipping for focal conv
-                    if fsize1 < fsize2:
-                        table_pretrained_resized = torch.zeros(table_current.shape)
-                        table_pretrained_resized[:, :, (fsize2-fsize1)//2:-(fsize2-fsize1)//2, (fsize2-fsize1)//2:-(fsize2-fsize1)//2] = table_pretrained
-                        v = table_pretrained_resized
-                    elif fsize1 > fsize2:
-                        table_pretrained_resized = table_pretrained[:, :, (fsize1-fsize2)//2:-(fsize1-fsize2)//2, (fsize1-fsize2)//2:-(fsize1-fsize2)//2]
-                        v = table_pretrained_resized
-
-
-                if ("modulation.f" in k or "pre_conv" in k): 
-                    table_pretrained = v
-                    table_current = model_dict[k]
-                    if table_pretrained.shape != table_current.shape:
-                        if len(table_pretrained.shape) == 2:
-                            dim = table_pretrained.shape[1]
-                            assert table_current.shape[1] == dim
-                            L1 = table_pretrained.shape[0]
-                            L2 = table_current.shape[0]
-
-                            if L1 < L2:
-                                table_pretrained_resized = torch.zeros(table_current.shape)
-                                # copy for linear project
-                                table_pretrained_resized[:2*dim] = table_pretrained[:2*dim]
-                                # copy for global token gating
-                                table_pretrained_resized[-1] = table_pretrained[-1]
-                                # copy for first multiple focal levels
-                                table_pretrained_resized[2*dim:2*dim+(L1-2*dim-1)] = table_pretrained[2*dim:-1]
-                                # reassign pretrained weights
-                                v = table_pretrained_resized
-                            elif L1 > L2:
-                                raise NotImplementedError
-                        elif len(table_pretrained.shape) == 1:
-                            dim = table_pretrained.shape[0]
-                            L1 = table_pretrained.shape[0]
-                            L2 = table_current.shape[0]
-                            if L1 < L2:
-                                table_pretrained_resized = torch.zeros(table_current.shape)
-                                # copy for linear project
-                                table_pretrained_resized[:dim] = table_pretrained[:dim]
-                                # copy for global token gating
-                                table_pretrained_resized[-1] = table_pretrained[-1]
-                                # copy for first multiple focal levels
-                                # table_pretrained_resized[dim:2*dim+(L1-2*dim-1)] = table_pretrained[2*dim:-1]
-                                # reassign pretrained weights
-                                v = table_pretrained_resized
-                            elif L1 > L2:
-                                raise NotImplementedError    
-
-                need_init_state_dict[k] = v
-        
-        self.load_state_dict(need_init_state_dict, strict=False)
-
-
-    def forward(self, x):
-        """Forward function."""
-        tic = time.time()
-        x = self.patch_embed(x)
-        Wh, Ww = x.size(2), x.size(3)
-
-        x = x.flatten(2).transpose(1, 2)
-        x = self.pos_drop(x)
-
-        outs = {}
-        for i in range(self.num_layers):
-            layer = self.layers[i]
-            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
-            if i in self.out_indices:
-                norm_layer = getattr(self, f'norm{i}')
-                x_out = norm_layer(x_out)
-
-                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
-                outs["res{}".format(i + 2)] = out
-                
-        if len(self.out_indices) == 0:
-            outs["res5"] = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
-
-        toc = time.time()
-        return outs
-
-    def train(self, mode=True):
-        """Convert the model into training mode while keep layers freezed."""
-        super(FocalNet, self).train(mode)
-        self._freeze_stages()
-
-
-class D2FocalNet(FocalNet, Backbone):
-    def __init__(self, cfg, input_shape):
-
-        pretrain_img_size = cfg['BACKBONE']['FOCAL']['PRETRAIN_IMG_SIZE']
-        patch_size = cfg['BACKBONE']['FOCAL']['PATCH_SIZE']
-        in_chans = 3
-        embed_dim = cfg['BACKBONE']['FOCAL']['EMBED_DIM']
-        depths = cfg['BACKBONE']['FOCAL']['DEPTHS']
-        mlp_ratio = cfg['BACKBONE']['FOCAL']['MLP_RATIO']
-        drop_rate = cfg['BACKBONE']['FOCAL']['DROP_RATE']
-        drop_path_rate = cfg['BACKBONE']['FOCAL']['DROP_PATH_RATE']
-        norm_layer = nn.LayerNorm
-        patch_norm = cfg['BACKBONE']['FOCAL']['PATCH_NORM']
-        use_checkpoint = cfg['BACKBONE']['FOCAL']['USE_CHECKPOINT']
-        out_indices = cfg['BACKBONE']['FOCAL']['OUT_INDICES']
-        scaling_modulator = cfg['BACKBONE']['FOCAL'].get('SCALING_MODULATOR', False)
-
-        super().__init__(
-            pretrain_img_size,
-            patch_size,
-            in_chans,
-            embed_dim,
-            depths,
-            mlp_ratio,
-            drop_rate,
-            drop_path_rate,
-            norm_layer,
-            patch_norm,
-            out_indices,
-            focal_levels=cfg['BACKBONE']['FOCAL']['FOCAL_LEVELS'],
-            focal_windows=cfg['BACKBONE']['FOCAL']['FOCAL_WINDOWS'],   
-            use_conv_embed=cfg['BACKBONE']['FOCAL']['USE_CONV_EMBED'],    
-            use_postln=cfg['BACKBONE']['FOCAL']['USE_POSTLN'],       
-            use_postln_in_modulation=cfg['BACKBONE']['FOCAL']['USE_POSTLN_IN_MODULATION'], 
-            scaling_modulator=scaling_modulator,
-            use_layerscale=cfg['BACKBONE']['FOCAL']['USE_LAYERSCALE'], 
-            use_checkpoint=use_checkpoint,
-        )
-
-        self._out_features = cfg['BACKBONE']['FOCAL']['OUT_FEATURES']
-
-        self._out_feature_strides = {
-            "res2": 4,
-            "res3": 8,
-            "res4": 16,
-            "res5": 32,
-        }
-        self._out_feature_channels = {
-            "res2": self.num_features[0],
-            "res3": self.num_features[1],
-            "res4": self.num_features[2],
-            "res5": self.num_features[3],
-        }
-
-    def forward(self, x):
-        """
-        Args:
-            x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
-        Returns:
-            dict[str->Tensor]: names and the corresponding features
-        """
-        assert (
-            x.dim() == 4
-        ), f"SwinTransformer takes an input of shape (N, C, H, W). Got {x.shape} instead!"
-        outputs = {}
-        y = super().forward(x)
-        for k in y.keys():
-            if k in self._out_features:
-                outputs[k] = y[k]
-        return outputs
-
-    def output_shape(self):
-        return {
-            name: ShapeSpec(
-                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
-            )
-            for name in self._out_features
-        }
-
-    @property
-    def size_divisibility(self):
-        return 32
-
-@register_backbone
-def get_focal_backbone(cfg):
-    focal = D2FocalNet(cfg['MODEL'], 224)    
-
-    if cfg['MODEL']['BACKBONE']['LOAD_PRETRAINED'] is True:
-        filename = cfg['MODEL']['BACKBONE']['PRETRAINED']
-        logger.info(f'=> init from {filename}')
-        with PathManager.open(filename, "rb") as f:
-            ckpt = torch.load(f)['model']
-        focal.load_weights(ckpt, cfg['MODEL']['BACKBONE']['FOCAL'].get('PRETRAINED_LAYERS', ['*']), cfg['VERBOSE'])
-
-    return focal

xdecoder/backbone/registry.py

@@ -1,14 +0,0 @@
-_model_entrypoints = {}
-
-
-def register_backbone(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/backbone/resnet.py

@@ -1,731 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-import pickle
-import numpy as np
-from typing import Any, Dict
-import fvcore.nn.weight_init as weight_init
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-
-from .backbone import Backbone
-from .registry import register_backbone
-
-from detectron2.layers import (
-    CNNBlockBase,
-    Conv2d,
-    DeformConv,
-    ModulatedDeformConv,
-    ShapeSpec,
-    get_norm,
-)
-from detectron2.utils.file_io import PathManager
-
-__all__ = [
-    "ResNetBlockBase",
-    "BasicBlock",
-    "BottleneckBlock",
-    "DeformBottleneckBlock",
-    "BasicStem",
-    "ResNet",
-    "make_stage",
-    "get_resnet_backbone",
-]
-
-
-class BasicBlock(CNNBlockBase):
-    """
-    The basic residual block for ResNet-18 and ResNet-34 defined in :paper:`ResNet`,
-    with two 3x3 conv layers and a projection shortcut if needed.
-    """
-
-    def __init__(self, in_channels, out_channels, *, stride=1, norm="BN"):
-        """
-        Args:
-            in_channels (int): Number of input channels.
-            out_channels (int): Number of output channels.
-            stride (int): Stride for the first conv.
-            norm (str or callable): normalization for all conv layers.
-                See :func:`layers.get_norm` for supported format.
-        """
-        super().__init__(in_channels, out_channels, stride)
-
-        if in_channels != out_channels:
-            self.shortcut = Conv2d(
-                in_channels,
-                out_channels,
-                kernel_size=1,
-                stride=stride,
-                bias=False,
-                norm=get_norm(norm, out_channels),
-            )
-        else:
-            self.shortcut = None
-
-        self.conv1 = Conv2d(
-            in_channels,
-            out_channels,
-            kernel_size=3,
-            stride=stride,
-            padding=1,
-            bias=False,
-            norm=get_norm(norm, out_channels),
-        )
-
-        self.conv2 = Conv2d(
-            out_channels,
-            out_channels,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            bias=False,
-            norm=get_norm(norm, out_channels),
-        )
-
-        for layer in [self.conv1, self.conv2, self.shortcut]:
-            if layer is not None:  # shortcut can be None
-                weight_init.c2_msra_fill(layer)
-
-    def forward(self, x):
-        out = self.conv1(x)
-        out = F.relu_(out)
-        out = self.conv2(out)
-
-        if self.shortcut is not None:
-            shortcut = self.shortcut(x)
-        else:
-            shortcut = x
-
-        out += shortcut
-        out = F.relu_(out)
-        return out
-
-
-class BottleneckBlock(CNNBlockBase):
-    """
-    The standard bottleneck residual block used by ResNet-50, 101 and 152
-    defined in :paper:`ResNet`.  It contains 3 conv layers with kernels
-    1x1, 3x3, 1x1, and a projection shortcut if needed.
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        *,
-        bottleneck_channels,
-        stride=1,
-        num_groups=1,
-        norm="BN",
-        stride_in_1x1=False,
-        dilation=1,
-    ):
-        """
-        Args:
-            bottleneck_channels (int): number of output channels for the 3x3
-                "bottleneck" conv layers.
-            num_groups (int): number of groups for the 3x3 conv layer.
-            norm (str or callable): normalization for all conv layers.
-                See :func:`layers.get_norm` for supported format.
-            stride_in_1x1 (bool): when stride>1, whether to put stride in the
-                first 1x1 convolution or the bottleneck 3x3 convolution.
-            dilation (int): the dilation rate of the 3x3 conv layer.
-        """
-        super().__init__(in_channels, out_channels, stride)
-
-        if in_channels != out_channels:
-            self.shortcut = Conv2d(
-                in_channels,
-                out_channels,
-                kernel_size=1,
-                stride=stride,
-                bias=False,
-                norm=get_norm(norm, out_channels),
-            )
-        else:
-            self.shortcut = None
-
-        # The original MSRA ResNet models have stride in the first 1x1 conv
-        # The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
-        # stride in the 3x3 conv
-        stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
-
-        self.conv1 = Conv2d(
-            in_channels,
-            bottleneck_channels,
-            kernel_size=1,
-            stride=stride_1x1,
-            bias=False,
-            norm=get_norm(norm, bottleneck_channels),
-        )
-
-        self.conv2 = Conv2d(
-            bottleneck_channels,
-            bottleneck_channels,
-            kernel_size=3,
-            stride=stride_3x3,
-            padding=1 * dilation,
-            bias=False,
-            groups=num_groups,
-            dilation=dilation,
-            norm=get_norm(norm, bottleneck_channels),
-        )
-
-        self.conv3 = Conv2d(
-            bottleneck_channels,
-            out_channels,
-            kernel_size=1,
-            bias=False,
-            norm=get_norm(norm, out_channels),
-        )
-
-        for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:
-            if layer is not None:  # shortcut can be None
-                weight_init.c2_msra_fill(layer)
-
-        # Zero-initialize the last normalization in each residual branch,
-        # so that at the beginning, the residual branch starts with zeros,
-        # and each residual block behaves like an identity.
-        # See Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
-        # "For BN layers, the learnable scaling coefficient γ is initialized
-        # to be 1, except for each residual block's last BN
-        # where γ is initialized to be 0."
-
-        # nn.init.constant_(self.conv3.norm.weight, 0)
-        # TODO this somehow hurts performance when training GN models from scratch.
-        # Add it as an option when we need to use this code to train a backbone.
-
-    def forward(self, x):
-        out = self.conv1(x)
-        out = F.relu_(out)
-
-        out = self.conv2(out)
-        out = F.relu_(out)
-
-        out = self.conv3(out)
-
-        if self.shortcut is not None:
-            shortcut = self.shortcut(x)
-        else:
-            shortcut = x
-
-        out += shortcut
-        out = F.relu_(out)
-        return out
-
-
-class DeformBottleneckBlock(CNNBlockBase):
-    """
-    Similar to :class:`BottleneckBlock`, but with :paper:`deformable conv <deformconv>`
-    in the 3x3 convolution.
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        *,
-        bottleneck_channels,
-        stride=1,
-        num_groups=1,
-        norm="BN",
-        stride_in_1x1=False,
-        dilation=1,
-        deform_modulated=False,
-        deform_num_groups=1,
-    ):
-        super().__init__(in_channels, out_channels, stride)
-        self.deform_modulated = deform_modulated
-
-        if in_channels != out_channels:
-            self.shortcut = Conv2d(
-                in_channels,
-                out_channels,
-                kernel_size=1,
-                stride=stride,
-                bias=False,
-                norm=get_norm(norm, out_channels),
-            )
-        else:
-            self.shortcut = None
-
-        stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
-
-        self.conv1 = Conv2d(
-            in_channels,
-            bottleneck_channels,
-            kernel_size=1,
-            stride=stride_1x1,
-            bias=False,
-            norm=get_norm(norm, bottleneck_channels),
-        )
-
-        if deform_modulated:
-            deform_conv_op = ModulatedDeformConv
-            # offset channels are 2 or 3 (if with modulated) * kernel_size * kernel_size
-            offset_channels = 27
-        else:
-            deform_conv_op = DeformConv
-            offset_channels = 18
-
-        self.conv2_offset = Conv2d(
-            bottleneck_channels,
-            offset_channels * deform_num_groups,
-            kernel_size=3,
-            stride=stride_3x3,
-            padding=1 * dilation,
-            dilation=dilation,
-        )
-        self.conv2 = deform_conv_op(
-            bottleneck_channels,
-            bottleneck_channels,
-            kernel_size=3,
-            stride=stride_3x3,
-            padding=1 * dilation,
-            bias=False,
-            groups=num_groups,
-            dilation=dilation,
-            deformable_groups=deform_num_groups,
-            norm=get_norm(norm, bottleneck_channels),
-        )
-
-        self.conv3 = Conv2d(
-            bottleneck_channels,
-            out_channels,
-            kernel_size=1,
-            bias=False,
-            norm=get_norm(norm, out_channels),
-        )
-
-        for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:
-            if layer is not None:  # shortcut can be None
-                weight_init.c2_msra_fill(layer)
-
-        nn.init.constant_(self.conv2_offset.weight, 0)
-        nn.init.constant_(self.conv2_offset.bias, 0)
-
-    def forward(self, x):
-        out = self.conv1(x)
-        out = F.relu_(out)
-
-        if self.deform_modulated:
-            offset_mask = self.conv2_offset(out)
-            offset_x, offset_y, mask = torch.chunk(offset_mask, 3, dim=1)
-            offset = torch.cat((offset_x, offset_y), dim=1)
-            mask = mask.sigmoid()
-            out = self.conv2(out, offset, mask)
-        else:
-            offset = self.conv2_offset(out)
-            out = self.conv2(out, offset)
-        out = F.relu_(out)
-
-        out = self.conv3(out)
-
-        if self.shortcut is not None:
-            shortcut = self.shortcut(x)
-        else:
-            shortcut = x
-
-        out += shortcut
-        out = F.relu_(out)
-        return out
-
-
-class BasicStem(CNNBlockBase):
-    """
-    The standard ResNet stem (layers before the first residual block),
-    with a conv, relu and max_pool.
-    """
-
-    def __init__(self, in_channels=3, out_channels=64, norm="BN"):
-        """
-        Args:
-            norm (str or callable): norm after the first conv layer.
-                See :func:`layers.get_norm` for supported format.
-        """
-        super().__init__(in_channels, out_channels, 4)
-        self.in_channels = in_channels
-        self.conv1 = Conv2d(
-            in_channels,
-            out_channels,
-            kernel_size=7,
-            stride=2,
-            padding=3,
-            bias=False,
-            norm=get_norm(norm, out_channels),
-        )
-        weight_init.c2_msra_fill(self.conv1)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = F.relu_(x)
-        x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
-        return x
-
-
-class ResNet(Backbone):
-    """
-    Implement :paper:`ResNet`.
-    """
-
-    def __init__(self, stem, stages, num_classes=None, out_features=None, freeze_at=0):
-        """
-        Args:
-            stem (nn.Module): a stem module
-            stages (list[list[CNNBlockBase]]): several (typically 4) stages,
-                each contains multiple :class:`CNNBlockBase`.
-            num_classes (None or int): if None, will not perform classification.
-                Otherwise, will create a linear layer.
-            out_features (list[str]): name of the layers whose outputs should
-                be returned in forward. Can be anything in "stem", "linear", or "res2" ...
-                If None, will return the output of the last layer.
-            freeze_at (int): The number of stages at the beginning to freeze.
-                see :meth:`freeze` for detailed explanation.
-        """
-        super().__init__()
-        self.stem = stem
-        self.num_classes = num_classes
-
-        current_stride = self.stem.stride
-        self._out_feature_strides = {"stem": current_stride}
-        self._out_feature_channels = {"stem": self.stem.out_channels}
-
-        self.stage_names, self.stages = [], []
-
-        if out_features is not None:
-            # Avoid keeping unused layers in this module. They consume extra memory
-            # and may cause allreduce to fail
-            num_stages = max(
-                [{"res2": 1, "res3": 2, "res4": 3, "res5": 4}.get(f, 0) for f in out_features]
-            )
-            stages = stages[:num_stages]
-        for i, blocks in enumerate(stages):
-            assert len(blocks) > 0, len(blocks)
-            for block in blocks:
-                assert isinstance(block, CNNBlockBase), block
-
-            name = "res" + str(i + 2)
-            stage = nn.Sequential(*blocks)
-
-            self.add_module(name, stage)
-            self.stage_names.append(name)
-            self.stages.append(stage)
-
-            self._out_feature_strides[name] = current_stride = int(
-                current_stride * np.prod([k.stride for k in blocks])
-            )
-            self._out_feature_channels[name] = curr_channels = blocks[-1].out_channels
-        self.stage_names = tuple(self.stage_names)  # Make it static for scripting
-
-        if num_classes is not None:
-            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
-            self.linear = nn.Linear(curr_channels, num_classes)
-
-            # Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
-            # "The 1000-way fully-connected layer is initialized by
-            # drawing weights from a zero-mean Gaussian with standard deviation of 0.01."
-            nn.init.normal_(self.linear.weight, std=0.01)
-            name = "linear"
-
-        if out_features is None:
-            out_features = [name]
-        self._out_features = out_features
-        assert len(self._out_features)
-        children = [x[0] for x in self.named_children()]
-        for out_feature in self._out_features:
-            assert out_feature in children, "Available children: {}".format(", ".join(children))
-        self.freeze(freeze_at)
-
-    def forward(self, x):
-        """
-        Args:
-            x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
-
-        Returns:
-            dict[str->Tensor]: names and the corresponding features
-        """
-        assert x.dim() == 4, f"ResNet takes an input of shape (N, C, H, W). Got {x.shape} instead!"
-        outputs = {}
-        x = self.stem(x)
-        if "stem" in self._out_features:
-            outputs["stem"] = x
-        for name, stage in zip(self.stage_names, self.stages):
-            x = stage(x)
-            if name in self._out_features:
-                outputs[name] = x
-        if self.num_classes is not None:
-            x = self.avgpool(x)
-            x = torch.flatten(x, 1)
-            x = self.linear(x)
-            if "linear" in self._out_features:
-                outputs["linear"] = x
-        return outputs
-
-    def output_shape(self):
-        return {
-            name: ShapeSpec(
-                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
-            )
-            for name in self._out_features
-        }
-
-    def freeze(self, freeze_at=0):
-        """
-        Freeze the first several stages of the ResNet. Commonly used in
-        fine-tuning.
-
-        Layers that produce the same feature map spatial size are defined as one
-        "stage" by :paper:`FPN`.
-
-        Args:
-            freeze_at (int): number of stages to freeze.
-                `1` means freezing the stem. `2` means freezing the stem and
-                one residual stage, etc.
-
-        Returns:
-            nn.Module: this ResNet itself
-        """
-        if freeze_at >= 1:
-            self.stem.freeze()
-        for idx, stage in enumerate(self.stages, start=2):
-            if freeze_at >= idx:
-                for block in stage.children():
-                    block.freeze()
-        return self
-
-    @staticmethod
-    def make_stage(block_class, num_blocks, *, in_channels, out_channels, **kwargs):
-        """
-        Create a list of blocks of the same type that forms one ResNet stage.
-
-        Args:
-            block_class (type): a subclass of CNNBlockBase that's used to create all blocks in this
-                stage. A module of this type must not change spatial resolution of inputs unless its
-                stride != 1.
-            num_blocks (int): number of blocks in this stage
-            in_channels (int): input channels of the entire stage.
-            out_channels (int): output channels of **every block** in the stage.
-            kwargs: other arguments passed to the constructor of
-                `block_class`. If the argument name is "xx_per_block", the
-                argument is a list of values to be passed to each block in the
-                stage. Otherwise, the same argument is passed to every block
-                in the stage.
-
-        Returns:
-            list[CNNBlockBase]: a list of block module.
-
-        Examples:
-        ::
-            stage = ResNet.make_stage(
-                BottleneckBlock, 3, in_channels=16, out_channels=64,
-                bottleneck_channels=16, num_groups=1,
-                stride_per_block=[2, 1, 1],
-                dilations_per_block=[1, 1, 2]
-            )
-
-        Usually, layers that produce the same feature map spatial size are defined as one
-        "stage" (in :paper:`FPN`). Under such definition, ``stride_per_block[1:]`` should
-        all be 1.
-        """
-        blocks = []
-        for i in range(num_blocks):
-            curr_kwargs = {}
-            for k, v in kwargs.items():
-                if k.endswith("_per_block"):
-                    assert len(v) == num_blocks, (
-                        f"Argument '{k}' of make_stage should have the "
-                        f"same length as num_blocks={num_blocks}."
-                    )
-                    newk = k[: -len("_per_block")]
-                    assert newk not in kwargs, f"Cannot call make_stage with both {k} and {newk}!"
-                    curr_kwargs[newk] = v[i]
-                else:
-                    curr_kwargs[k] = v
-
-            blocks.append(
-                block_class(in_channels=in_channels, out_channels=out_channels, **curr_kwargs)
-            )
-            in_channels = out_channels
-        return blocks
-
-    @staticmethod
-    def make_default_stages(depth, block_class=None, **kwargs):
-        """
-        Created list of ResNet stages from pre-defined depth (one of 18, 34, 50, 101, 152).
-        If it doesn't create the ResNet variant you need, please use :meth:`make_stage`
-        instead for fine-grained customization.
-
-        Args:
-            depth (int): depth of ResNet
-            block_class (type): the CNN block class. Has to accept
-                `bottleneck_channels` argument for depth > 50.
-                By default it is BasicBlock or BottleneckBlock, based on the
-                depth.
-            kwargs:
-                other arguments to pass to `make_stage`. Should not contain
-                stride and channels, as they are predefined for each depth.
-
-        Returns:
-            list[list[CNNBlockBase]]: modules in all stages; see arguments of
-                :class:`ResNet.__init__`.
-        """
-        num_blocks_per_stage = {
-            18: [2, 2, 2, 2],
-            34: [3, 4, 6, 3],
-            50: [3, 4, 6, 3],
-            101: [3, 4, 23, 3],
-            152: [3, 8, 36, 3],
-        }[depth]
-        if block_class is None:
-            block_class = BasicBlock if depth < 50 else BottleneckBlock
-        if depth < 50:
-            in_channels = [64, 64, 128, 256]
-            out_channels = [64, 128, 256, 512]
-        else:
-            in_channels = [64, 256, 512, 1024]
-            out_channels = [256, 512, 1024, 2048]
-        ret = []
-        for (n, s, i, o) in zip(num_blocks_per_stage, [1, 2, 2, 2], in_channels, out_channels):
-            if depth >= 50:
-                kwargs["bottleneck_channels"] = o // 4
-            ret.append(
-                ResNet.make_stage(
-                    block_class=block_class,
-                    num_blocks=n,
-                    stride_per_block=[s] + [1] * (n - 1),
-                    in_channels=i,
-                    out_channels=o,
-                    **kwargs,
-                )
-            )
-        return ret
-
-
-ResNetBlockBase = CNNBlockBase
-"""
-Alias for backward compatibiltiy.
-"""
-
-
-def make_stage(*args, **kwargs):
-    """
-    Deprecated alias for backward compatibiltiy.
-    """
-    return ResNet.make_stage(*args, **kwargs)
-
-
-def _convert_ndarray_to_tensor(state_dict: Dict[str, Any]) -> None:
-    """
-    In-place convert all numpy arrays in the state_dict to torch tensor.
-    Args:
-        state_dict (dict): a state-dict to be loaded to the model.
-            Will be modified.
-    """
-    # model could be an OrderedDict with _metadata attribute
-    # (as returned by Pytorch's state_dict()). We should preserve these
-    # properties.
-    for k in list(state_dict.keys()):
-        v = state_dict[k]
-        if not isinstance(v, np.ndarray) and not isinstance(v, torch.Tensor):
-            raise ValueError(
-                "Unsupported type found in checkpoint! {}: {}".format(k, type(v))
-            )
-        if not isinstance(v, torch.Tensor):
-            state_dict[k] = torch.from_numpy(v)
-
-
-@register_backbone
-def get_resnet_backbone(cfg):
-    """
-    Create a ResNet instance from config.
-
-    Returns:
-        ResNet: a :class:`ResNet` instance.
-    """
-    res_cfg = cfg['MODEL']['BACKBONE']['RESNETS']
-
-    # need registration of new blocks/stems?
-    norm = res_cfg['NORM']
-    stem = BasicStem(
-        in_channels=res_cfg['STEM_IN_CHANNELS'],
-        out_channels=res_cfg['STEM_OUT_CHANNELS'],
-        norm=norm,
-    )
-
-    # fmt: off
-    freeze_at           = res_cfg['FREEZE_AT']
-    out_features        = res_cfg['OUT_FEATURES']
-    depth               = res_cfg['DEPTH']
-    num_groups          = res_cfg['NUM_GROUPS']
-    width_per_group     = res_cfg['WIDTH_PER_GROUP']
-    bottleneck_channels = num_groups * width_per_group
-    in_channels         = res_cfg['STEM_OUT_CHANNELS']
-    out_channels        = res_cfg['RES2_OUT_CHANNELS']
-    stride_in_1x1       = res_cfg['STRIDE_IN_1X1']
-    res5_dilation       = res_cfg['RES5_DILATION']
-    deform_on_per_stage = res_cfg['DEFORM_ON_PER_STAGE']
-    deform_modulated    = res_cfg['DEFORM_MODULATED']
-    deform_num_groups   = res_cfg['DEFORM_NUM_GROUPS']
-    # fmt: on
-    assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)
-
-    num_blocks_per_stage = {
-        18: [2, 2, 2, 2],
-        34: [3, 4, 6, 3],
-        50: [3, 4, 6, 3],
-        101: [3, 4, 23, 3],
-        152: [3, 8, 36, 3],
-    }[depth]
-
-    if depth in [18, 34]:
-        assert out_channels == 64, "Must set MODEL.RESNETS.RES2_OUT_CHANNELS = 64 for R18/R34"
-        assert not any(
-            deform_on_per_stage
-        ), "MODEL.RESNETS.DEFORM_ON_PER_STAGE unsupported for R18/R34"
-        assert res5_dilation == 1, "Must set MODEL.RESNETS.RES5_DILATION = 1 for R18/R34"
-        assert num_groups == 1, "Must set MODEL.RESNETS.NUM_GROUPS = 1 for R18/R34"
-
-    stages = []
-
-    for idx, stage_idx in enumerate(range(2, 6)):
-        # res5_dilation is used this way as a convention in R-FCN & Deformable Conv paper
-        dilation = res5_dilation if stage_idx == 5 else 1
-        first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
-        stage_kargs = {
-            "num_blocks": num_blocks_per_stage[idx],
-            "stride_per_block": [first_stride] + [1] * (num_blocks_per_stage[idx] - 1),
-            "in_channels": in_channels,
-            "out_channels": out_channels,
-            "norm": norm,
-        }
-        # Use BasicBlock for R18 and R34.
-        if depth in [18, 34]:
-            stage_kargs["block_class"] = BasicBlock
-        else:
-            stage_kargs["bottleneck_channels"] = bottleneck_channels
-            stage_kargs["stride_in_1x1"] = stride_in_1x1
-            stage_kargs["dilation"] = dilation
-            stage_kargs["num_groups"] = num_groups
-            if deform_on_per_stage[idx]:
-                stage_kargs["block_class"] = DeformBottleneckBlock
-                stage_kargs["deform_modulated"] = deform_modulated
-                stage_kargs["deform_num_groups"] = deform_num_groups
-            else:
-                stage_kargs["block_class"] = BottleneckBlock
-        blocks = ResNet.make_stage(**stage_kargs)
-        in_channels = out_channels
-        out_channels *= 2
-        bottleneck_channels *= 2
-        stages.append(blocks)
-    backbone = ResNet(stem, stages, out_features=out_features, freeze_at=freeze_at)
-
-    if cfg['MODEL']['BACKBONE']['LOAD_PRETRAINED'] is True:
-        filename = cfg['MODEL']['BACKBONE']['PRETRAINED']
-        with PathManager.open(filename, "rb") as f:
-            ckpt = pickle.load(f, encoding="latin1")['model']
-        _convert_ndarray_to_tensor(ckpt)
-        ckpt.pop('stem.fc.weight')
-        ckpt.pop('stem.fc.bias')
-        backbone.load_state_dict(ckpt)
-
-    return backbone

xdecoder/backbone/swin.py

@@ -1,892 +0,0 @@
-# --------------------------------------------------------
-# Swin Transformer
-# Copyright (c) 2021 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Ze Liu, Yutong Lin, Yixuan Wei
-# --------------------------------------------------------
-
-# Copyright (c) Facebook, Inc. and its affiliates.
-# Modified by Bowen Cheng from https://github.com/SwinTransformer/Swin-Transformer-Semantic-Segmentation/blob/main/mmseg/models/backbones/swin_transformer.py
-import logging
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint as checkpoint
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
-
-from detectron2.modeling import Backbone, ShapeSpec
-from detectron2.utils.file_io import PathManager
-
-from .registry import register_backbone
-
-logger = logging.getLogger(__name__)
-
-
-class Mlp(nn.Module):
-    """Multilayer perceptron."""
-
-    def __init__(
-        self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0
-    ):
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(hidden_features, out_features)
-        self.drop = nn.Dropout(drop)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.drop(x)
-        x = self.fc2(x)
-        x = self.drop(x)
-        return x
-
-
-def window_partition(x, window_size):
-    """
-    Args:
-        x: (B, H, W, C)
-        window_size (int): window size
-    Returns:
-        windows: (num_windows*B, window_size, window_size, C)
-    """
-    B, H, W, C = x.shape
-    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
-    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
-    return windows
-
-
-def window_reverse(windows, window_size, H, W):
-    """
-    Args:
-        windows: (num_windows*B, window_size, window_size, C)
-        window_size (int): Window size
-        H (int): Height of image
-        W (int): Width of image
-    Returns:
-        x: (B, H, W, C)
-    """
-    B = int(windows.shape[0] / (H * W / window_size / window_size))
-    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
-    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
-    return x
-
-
-class WindowAttention(nn.Module):
-    """Window based multi-head self attention (W-MSA) module with relative position bias.
-    It supports both of shifted and non-shifted window.
-    Args:
-        dim (int): Number of input channels.
-        window_size (tuple[int]): The height and width of the window.
-        num_heads (int): Number of attention heads.
-        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
-        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
-        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
-    """
-
-    def __init__(
-        self,
-        dim,
-        window_size,
-        num_heads,
-        qkv_bias=True,
-        qk_scale=None,
-        attn_drop=0.0,
-        proj_drop=0.0,
-    ):
-
-        super().__init__()
-        self.dim = dim
-        self.window_size = window_size  # Wh, Ww
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        self.scale = qk_scale or head_dim ** -0.5
-
-        # define a parameter table of relative position bias
-        self.relative_position_bias_table = nn.Parameter(
-            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
-        )  # 2*Wh-1 * 2*Ww-1, nH
-
-        # get pair-wise relative position index for each token inside the window
-        coords_h = torch.arange(self.window_size[0])
-        coords_w = torch.arange(self.window_size[1])
-        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
-        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
-        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
-        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
-        relative_coords[:, :, 1] += self.window_size[1] - 1
-        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
-        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
-        self.register_buffer("relative_position_index", relative_position_index)
-
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.proj = nn.Linear(dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-
-        trunc_normal_(self.relative_position_bias_table, std=0.02)
-        self.softmax = nn.Softmax(dim=-1)
-
-    def forward(self, x, mask=None):
-        """Forward function.
-        Args:
-            x: input features with shape of (num_windows*B, N, C)
-            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
-        """
-        B_, N, C = x.shape
-        qkv = (
-            self.qkv(x)
-            .reshape(B_, N, 3, self.num_heads, C // self.num_heads)
-            .permute(2, 0, 3, 1, 4)
-        )
-        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
-
-        q = q * self.scale
-        attn = q @ k.transpose(-2, -1)
-        
-        relative_position_bias = self.relative_position_bias_table[
-            self.relative_position_index.view(-1)
-        ].view(
-            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
-        )  # Wh*Ww,Wh*Ww,nH
-        relative_position_bias = relative_position_bias.permute(
-            2, 0, 1
-        ).contiguous()  # nH, Wh*Ww, Wh*Ww
-        attn = attn + relative_position_bias.unsqueeze(0)
-
-        if mask is not None:
-            nW = mask.shape[0]
-            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
-            attn = attn.view(-1, self.num_heads, N, N)
-            attn = self.softmax(attn)
-        else:
-            attn = self.softmax(attn)
-
-        attn = self.attn_drop(attn)
-
-        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
-        x = self.proj(x)
-        x = self.proj_drop(x)
-
-        return x
-
-
-class SwinTransformerBlock(nn.Module):
-    """Swin Transformer Block.
-    Args:
-        dim (int): Number of input channels.
-        num_heads (int): Number of attention heads.
-        window_size (int): Window size.
-        shift_size (int): Shift size for SW-MSA.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
-        drop (float, optional): Dropout rate. Default: 0.0
-        attn_drop (float, optional): Attention dropout rate. Default: 0.0
-        drop_path (float, optional): Stochastic depth rate. Default: 0.0
-        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
-    """
-
-    def __init__(
-        self,
-        dim,
-        num_heads,
-        window_size=7,
-        shift_size=0,
-        mlp_ratio=4.0,
-        qkv_bias=True,
-        qk_scale=None,
-        drop=0.0,
-        attn_drop=0.0,
-        drop_path=0.0,
-        act_layer=nn.GELU,
-        norm_layer=nn.LayerNorm,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.num_heads = num_heads
-        self.window_size = window_size
-        self.shift_size = shift_size
-        self.mlp_ratio = mlp_ratio
-        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
-
-        self.norm1 = norm_layer(dim)
-        self.attn = WindowAttention(
-            dim,
-            window_size=to_2tuple(self.window_size),
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            qk_scale=qk_scale,
-            attn_drop=attn_drop,
-            proj_drop=drop,
-        )
-
-        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = Mlp(
-            in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop
-        )
-
-        self.H = None
-        self.W = None
-
-    def forward(self, x, mask_matrix):
-        """Forward function.
-        Args:
-            x: Input feature, tensor size (B, H*W, C).
-            H, W: Spatial resolution of the input feature.
-            mask_matrix: Attention mask for cyclic shift.
-        """
-        B, L, C = x.shape
-        H, W = self.H, self.W
-        assert L == H * W, "input feature has wrong size"
-
-        # HACK model will not upsampling
-        # if min([H, W]) <= self.window_size:
-            # if window size is larger than input resolution, we don't partition windows
-            # self.shift_size = 0
-            # self.window_size = min([H,W])
-
-        shortcut = x
-        x = self.norm1(x)
-        x = x.view(B, H, W, C)
-
-        # pad feature maps to multiples of window size
-        pad_l = pad_t = 0
-        pad_r = (self.window_size - W % self.window_size) % self.window_size
-        pad_b = (self.window_size - H % self.window_size) % self.window_size
-        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
-        _, Hp, Wp, _ = x.shape
-
-        # cyclic shift
-        if self.shift_size > 0:
-            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
-            attn_mask = mask_matrix
-        else:
-            shifted_x = x
-            attn_mask = None
-
-        # partition windows
-        x_windows = window_partition(
-            shifted_x, self.window_size
-        )  # nW*B, window_size, window_size, C
-        x_windows = x_windows.view(
-            -1, self.window_size * self.window_size, C
-        )  # nW*B, window_size*window_size, C
-
-        # W-MSA/SW-MSA
-        attn_windows = self.attn(x_windows, mask=attn_mask)  # nW*B, window_size*window_size, C
-
-        # merge windows
-        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
-        shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # B H' W' C
-
-        # reverse cyclic shift
-        if self.shift_size > 0:
-            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
-        else:
-            x = shifted_x
-
-        if pad_r > 0 or pad_b > 0:
-            x = x[:, :H, :W, :].contiguous()
-
-        x = x.view(B, H * W, C)
-
-        # FFN
-        x = shortcut + self.drop_path(x)
-        x = x + self.drop_path(self.mlp(self.norm2(x)))
-        return x
-
-
-class PatchMerging(nn.Module):
-    """Patch Merging Layer
-    Args:
-        dim (int): Number of input channels.
-        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
-    """
-
-    def __init__(self, dim, norm_layer=nn.LayerNorm):
-        super().__init__()
-        self.dim = dim
-        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
-        self.norm = norm_layer(4 * dim)
-
-    def forward(self, x, H, W):
-        """Forward function.
-        Args:
-            x: Input feature, tensor size (B, H*W, C).
-            H, W: Spatial resolution of the input feature.
-        """
-        B, L, C = x.shape
-        assert L == H * W, "input feature has wrong size"
-
-        x = x.view(B, H, W, C)
-
-        # padding
-        pad_input = (H % 2 == 1) or (W % 2 == 1)
-        if pad_input:
-            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
-
-        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
-        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
-        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
-        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
-        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
-        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
-
-        x = self.norm(x)
-        x = self.reduction(x)
-
-        return x
-
-
-class BasicLayer(nn.Module):
-    """A basic Swin Transformer layer for one stage.
-    Args:
-        dim (int): Number of feature channels
-        depth (int): Depths of this stage.
-        num_heads (int): Number of attention head.
-        window_size (int): Local window size. Default: 7.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
-        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
-        drop (float, optional): Dropout rate. Default: 0.0
-        attn_drop (float, optional): Attention dropout rate. Default: 0.0
-        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
-        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
-        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
-    """
-
-    def __init__(
-        self,
-        dim,
-        depth,
-        num_heads,
-        window_size=7,
-        mlp_ratio=4.0,
-        qkv_bias=True,
-        qk_scale=None,
-        drop=0.0,
-        attn_drop=0.0,
-        drop_path=0.0,
-        norm_layer=nn.LayerNorm,
-        downsample=None,
-        use_checkpoint=False,
-    ):
-        super().__init__()
-        self.window_size = window_size
-        self.shift_size = window_size // 2
-        self.depth = depth
-        self.use_checkpoint = use_checkpoint
-
-        # build blocks
-        self.blocks = nn.ModuleList(
-            [
-                SwinTransformerBlock(
-                    dim=dim,
-                    num_heads=num_heads,
-                    window_size=window_size,
-                    shift_size=0 if (i % 2 == 0) else window_size // 2,
-                    mlp_ratio=mlp_ratio,
-                    qkv_bias=qkv_bias,
-                    qk_scale=qk_scale,
-                    drop=drop,
-                    attn_drop=attn_drop,
-                    drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
-                    norm_layer=norm_layer,
-                )
-                for i in range(depth)
-            ]
-        )
-
-        # patch merging layer
-        if downsample is not None:
-            self.downsample = downsample(dim=dim, norm_layer=norm_layer)
-        else:
-            self.downsample = None
-
-    def forward(self, x, H, W):
-        """Forward function.
-        Args:
-            x: Input feature, tensor size (B, H*W, C).
-            H, W: Spatial resolution of the input feature.
-        """
-
-        # calculate attention mask for SW-MSA
-        Hp = int(np.ceil(H / self.window_size)) * self.window_size
-        Wp = int(np.ceil(W / self.window_size)) * self.window_size
-        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # 1 Hp Wp 1
-        h_slices = (
-            slice(0, -self.window_size),
-            slice(-self.window_size, -self.shift_size),
-            slice(-self.shift_size, None),
-        )
-        w_slices = (
-            slice(0, -self.window_size),
-            slice(-self.window_size, -self.shift_size),
-            slice(-self.shift_size, None),
-        )
-        cnt = 0
-        for h in h_slices:
-            for w in w_slices:
-                img_mask[:, h, w, :] = cnt
-                cnt += 1
-
-        mask_windows = window_partition(
-            img_mask, self.window_size
-        )  # nW, window_size, window_size, 1
-        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
-        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
-        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
-            attn_mask == 0, float(0.0)
-        ).type(x.dtype)
-        
-        for blk in self.blocks:
-            blk.H, blk.W = H, W
-            if self.use_checkpoint:
-                x = checkpoint.checkpoint(blk, x, attn_mask)
-            else:
-                x = blk(x, attn_mask)
-        if self.downsample is not None:
-            x_down = self.downsample(x, H, W)
-            Wh, Ww = (H + 1) // 2, (W + 1) // 2
-            return x, H, W, x_down, Wh, Ww
-        else:
-            return x, H, W, x, H, W
-
-
-class PatchEmbed(nn.Module):
-    """Image to Patch Embedding
-    Args:
-        patch_size (int): Patch token size. Default: 4.
-        in_chans (int): Number of input image channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        norm_layer (nn.Module, optional): Normalization layer. Default: None
-    """
-
-    def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
-        super().__init__()
-        patch_size = to_2tuple(patch_size)
-        self.patch_size = patch_size
-
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
-
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
-        if norm_layer is not None:
-            self.norm = norm_layer(embed_dim)
-        else:
-            self.norm = None
-
-    def forward(self, x):
-        """Forward function."""
-        # padding
-        _, _, H, W = x.size()
-        if W % self.patch_size[1] != 0:
-            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
-        if H % self.patch_size[0] != 0:
-            x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
-
-        x = self.proj(x)  # B C Wh Ww
-        if self.norm is not None:
-            Wh, Ww = x.size(2), x.size(3)
-            x = x.flatten(2).transpose(1, 2)
-            x = self.norm(x)
-            x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
-
-        return x
-
-
-class SwinTransformer(nn.Module):
-    """Swin Transformer backbone.
-        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
-          https://arxiv.org/pdf/2103.14030
-    Args:
-        pretrain_img_size (int): Input image size for training the pretrained model,
-            used in absolute postion embedding. Default 224.
-        patch_size (int | tuple(int)): Patch size. Default: 4.
-        in_chans (int): Number of input image channels. Default: 3.
-        embed_dim (int): Number of linear projection output channels. Default: 96.
-        depths (tuple[int]): Depths of each Swin Transformer stage.
-        num_heads (tuple[int]): Number of attention head of each stage.
-        window_size (int): Window size. Default: 7.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
-        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
-        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
-        drop_rate (float): Dropout rate.
-        attn_drop_rate (float): Attention dropout rate. Default: 0.
-        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
-        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
-        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False.
-        patch_norm (bool): If True, add normalization after patch embedding. Default: True.
-        out_indices (Sequence[int]): Output from which stages.
-        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
-            -1 means not freezing any parameters.
-        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
-    """
-
-    def __init__(
-        self,
-        pretrain_img_size=224,
-        patch_size=4,
-        in_chans=3,
-        embed_dim=96,
-        depths=[2, 2, 6, 2],
-        num_heads=[3, 6, 12, 24],
-        window_size=7,
-        mlp_ratio=4.0,
-        qkv_bias=True,
-        qk_scale=None,
-        drop_rate=0.0,
-        attn_drop_rate=0.0,
-        drop_path_rate=0.2,
-        norm_layer=nn.LayerNorm,
-        ape=False,
-        patch_norm=True,
-        out_indices=(0, 1, 2, 3),
-        frozen_stages=-1,
-        use_checkpoint=False,
-    ):
-        super().__init__()
-
-        self.pretrain_img_size = pretrain_img_size
-        self.num_layers = len(depths)
-        self.embed_dim = embed_dim
-        self.ape = ape
-        self.patch_norm = patch_norm
-        self.out_indices = out_indices
-        self.frozen_stages = frozen_stages
-
-        # split image into non-overlapping patches
-        self.patch_embed = PatchEmbed(
-            patch_size=patch_size,
-            in_chans=in_chans,
-            embed_dim=embed_dim,
-            norm_layer=norm_layer if self.patch_norm else None,
-        )
-
-        # absolute position embedding
-        if self.ape:
-            pretrain_img_size = to_2tuple(pretrain_img_size)
-            patch_size = to_2tuple(patch_size)
-            patches_resolution = [
-                pretrain_img_size[0] // patch_size[0],
-                pretrain_img_size[1] // patch_size[1],
-            ]
-
-            self.absolute_pos_embed = nn.Parameter(
-                torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])
-            )
-            trunc_normal_(self.absolute_pos_embed, std=0.02)
-
-        self.pos_drop = nn.Dropout(p=drop_rate)
-
-        # stochastic depth
-        dpr = [
-            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
-        ]  # stochastic depth decay rule
-
-        # build layers
-        self.layers = nn.ModuleList()
-        for i_layer in range(self.num_layers):
-            layer = BasicLayer(
-                dim=int(embed_dim * 2 ** i_layer),
-                depth=depths[i_layer],
-                num_heads=num_heads[i_layer],
-                window_size=window_size,
-                mlp_ratio=mlp_ratio,
-                qkv_bias=qkv_bias,
-                qk_scale=qk_scale,
-                drop=drop_rate,
-                attn_drop=attn_drop_rate,
-                drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],
-                norm_layer=norm_layer,
-                downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
-                use_checkpoint=use_checkpoint,
-            )
-            self.layers.append(layer)
-
-        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
-        self.num_features = num_features
-
-        # add a norm layer for each output
-        for i_layer in out_indices:
-            layer = norm_layer(num_features[i_layer])
-            layer_name = f"norm{i_layer}"
-            self.add_module(layer_name, layer)
-
-        self._freeze_stages()
-
-    def _freeze_stages(self):
-        if self.frozen_stages >= 0:
-            self.patch_embed.eval()
-            for param in self.patch_embed.parameters():
-                param.requires_grad = False
-
-        if self.frozen_stages >= 1 and self.ape:
-            self.absolute_pos_embed.requires_grad = False
-
-        if self.frozen_stages >= 2:
-            self.pos_drop.eval()
-            for i in range(0, self.frozen_stages - 1):
-                m = self.layers[i]
-                m.eval()
-                for param in m.parameters():
-                    param.requires_grad = False
-
-    def init_weights(self, pretrained=None):
-        """Initialize the weights in backbone.
-        Args:
-            pretrained (str, optional): Path to pre-trained weights.
-                Defaults to None.
-        """
-
-        def _init_weights(m):
-            if isinstance(m, nn.Linear):
-                trunc_normal_(m.weight, std=0.02)
-                if isinstance(m, nn.Linear) and m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.LayerNorm):
-                nn.init.constant_(m.bias, 0)
-                nn.init.constant_(m.weight, 1.0)
-
-
-    def load_weights(self, pretrained_dict=None, pretrained_layers=[], verbose=True):
-        model_dict = self.state_dict()
-        pretrained_dict = {
-            k: v for k, v in pretrained_dict.items()
-            if k in model_dict.keys()
-        }
-        need_init_state_dict = {}
-        for k, v in pretrained_dict.items():
-            need_init = (
-                    (
-                            k.split('.')[0] in pretrained_layers
-                            or pretrained_layers[0] == '*'
-                    )
-                    and 'relative_position_index' not in k
-                    and 'attn_mask' not in k
-            )
-
-            if need_init:
-                # if verbose:
-                #     logger.info(f'=> init {k} from {pretrained}')
-
-                if 'relative_position_bias_table' in k and v.size() != model_dict[k].size():
-                    relative_position_bias_table_pretrained = v
-                    relative_position_bias_table_current = model_dict[k]
-                    L1, nH1 = relative_position_bias_table_pretrained.size()
-                    L2, nH2 = relative_position_bias_table_current.size()
-                    if nH1 != nH2:
-                        logger.info(f"Error in loading {k}, passing")
-                    else:
-                        if L1 != L2:
-                            logger.info(
-                                '=> load_pretrained: resized variant: {} to {}'
-                                    .format((L1, nH1), (L2, nH2))
-                            )
-                            S1 = int(L1 ** 0.5)
-                            S2 = int(L2 ** 0.5)
-                            relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate(
-                                relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1),
-                                size=(S2, S2),
-                                mode='bicubic')
-                            v = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0)
-
-                if 'absolute_pos_embed' in k and v.size() != model_dict[k].size():
-                    absolute_pos_embed_pretrained = v
-                    absolute_pos_embed_current = model_dict[k]
-                    _, L1, C1 = absolute_pos_embed_pretrained.size()
-                    _, L2, C2 = absolute_pos_embed_current.size()
-                    if C1 != C1:
-                        logger.info(f"Error in loading {k}, passing")
-                    else:
-                        if L1 != L2:
-                            logger.info(
-                                '=> load_pretrained: resized variant: {} to {}'
-                                    .format((1, L1, C1), (1, L2, C2))
-                            )
-                            S1 = int(L1 ** 0.5)
-                            S2 = int(L2 ** 0.5)
-                            absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape(-1, S1, S1, C1)
-                            absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute(0, 3, 1, 2)
-                            absolute_pos_embed_pretrained_resized = torch.nn.functional.interpolate(
-                                absolute_pos_embed_pretrained, size=(S2, S2), mode='bicubic')
-                            v = absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1).flatten(1, 2)
-
-                need_init_state_dict[k] = v
-        self.load_state_dict(need_init_state_dict, strict=False)
-
-
-    def forward(self, x):
-        """Forward function."""
-        x = self.patch_embed(x)
-
-        Wh, Ww = x.size(2), x.size(3)
-        if self.ape:
-            # interpolate the position embedding to the corresponding size
-            absolute_pos_embed = F.interpolate(
-                self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic"
-            )
-            x = (x + absolute_pos_embed).flatten(2).transpose(1, 2)  # B Wh*Ww C
-        else:
-            x = x.flatten(2).transpose(1, 2)
-        x = self.pos_drop(x)
-
-        outs = {}
-        for i in range(self.num_layers):
-            layer = self.layers[i]
-            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
-
-            if i in self.out_indices:
-                norm_layer = getattr(self, f"norm{i}")
-                x_out = norm_layer(x_out)
-
-                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
-                outs["res{}".format(i + 2)] = out
-
-        if len(self.out_indices) == 0:
-            outs["res5"] = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
-        
-
-        return outs
-
-    def train(self, mode=True):
-        """Convert the model into training mode while keep layers freezed."""
-        super(SwinTransformer, self).train(mode)
-        self._freeze_stages()
-
-
-class D2SwinTransformer(SwinTransformer, Backbone):
-    def __init__(self, cfg, pretrain_img_size, patch_size, in_chans, embed_dim, 
-                 depths, num_heads, window_size, mlp_ratio, qkv_bias, qk_scale,
-                 drop_rate, attn_drop_rate, drop_path_rate, norm_layer, ape, 
-                 patch_norm, out_indices, use_checkpoint):
-        super().__init__(
-            pretrain_img_size,
-            patch_size,
-            in_chans,
-            embed_dim,
-            depths,
-            num_heads,
-            window_size,
-            mlp_ratio,
-            qkv_bias,
-            qk_scale,
-            drop_rate,
-            attn_drop_rate,
-            drop_path_rate,
-            norm_layer,
-            ape,
-            patch_norm,
-            out_indices,
-            use_checkpoint=use_checkpoint,
-        )
-
-        self._out_features = cfg['OUT_FEATURES']
-
-        self._out_feature_strides = {
-            "res2": 4,
-            "res3": 8,
-            "res4": 16,
-            "res5": 32,
-        }
-        self._out_feature_channels = {
-            "res2": self.num_features[0],
-            "res3": self.num_features[1],
-            "res4": self.num_features[2],
-            "res5": self.num_features[3],
-        }
-
-    def forward(self, x):
-        """
-        Args:
-            x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
-        Returns:
-            dict[str->Tensor]: names and the corresponding features
-        """
-        assert (
-            x.dim() == 4
-        ), f"SwinTransformer takes an input of shape (N, C, H, W). Got {x.shape} instead!"
-        outputs = {}
-        y = super().forward(x)
-        for k in y.keys():
-            if k in self._out_features:
-                outputs[k] = y[k]
-        return outputs
-
-    def output_shape(self):
-        feature_names = list(set(self._out_feature_strides.keys()) & set(self._out_features))
-        return {
-            name: ShapeSpec(
-                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
-            )
-            for name in feature_names
-        }
-
-    @property
-    def size_divisibility(self):
-        return 32
-
-
-@register_backbone
-def get_swin_backbone(cfg):
-    swin_cfg = cfg['MODEL']['BACKBONE']['SWIN']
-
-    pretrain_img_size = swin_cfg['PRETRAIN_IMG_SIZE']
-    patch_size = swin_cfg['PATCH_SIZE']
-    in_chans = 3
-    embed_dim = swin_cfg['EMBED_DIM']
-    depths = swin_cfg['DEPTHS']
-    num_heads = swin_cfg['NUM_HEADS']
-    window_size = swin_cfg['WINDOW_SIZE']
-    mlp_ratio = swin_cfg['MLP_RATIO']
-    qkv_bias = swin_cfg['QKV_BIAS']
-    qk_scale = swin_cfg['QK_SCALE']
-    drop_rate = swin_cfg['DROP_RATE']
-    attn_drop_rate = swin_cfg['ATTN_DROP_RATE']
-    drop_path_rate = swin_cfg['DROP_PATH_RATE']
-    norm_layer = nn.LayerNorm
-    ape = swin_cfg['APE']
-    patch_norm = swin_cfg['PATCH_NORM']
-    use_checkpoint = swin_cfg['USE_CHECKPOINT']
-    out_indices = swin_cfg.get('OUT_INDICES', [0,1,2,3])
-    
-    swin = D2SwinTransformer(
-        swin_cfg,
-        pretrain_img_size,
-        patch_size,
-        in_chans,
-        embed_dim,
-        depths,
-        num_heads,
-        window_size,
-        mlp_ratio,
-        qkv_bias,
-        qk_scale,
-        drop_rate,
-        attn_drop_rate,
-        drop_path_rate,
-        norm_layer,
-        ape,
-        patch_norm,
-        out_indices,
-        use_checkpoint=use_checkpoint,
-    )    
-
-    if cfg['MODEL']['BACKBONE']['LOAD_PRETRAINED'] is True:
-        filename = cfg['MODEL']['BACKBONE']['PRETRAINED']
-        with PathManager.open(filename, "rb") as f:
-            ckpt = torch.load(f, map_location=cfg['device'])['model']
-        swin.load_weights(ckpt, swin_cfg.get('PRETRAINED_LAYERS', ['*']), cfg['VERBOSE'])
-
-    return swin

xdecoder/body/__init__.py

@@ -1 +0,0 @@
-from .build import build_xdecoder_head

xdecoder/body/build.py

@@ -1,13 +0,0 @@
-from .registry import model_entrypoints
-from .registry import is_model
-
-from .xdecoder_head import *
-
-
-def build_xdecoder_head(config, *args, **kwargs):
-    model_name = config['MODEL']['HEAD']
-    if not is_model(model_name):
-        raise ValueError(f'Unkown model: {model_name}')
-
-    body = model_entrypoints(model_name)(config, *args, **kwargs)
-    return body

xdecoder/body/decoder/__init__.py

@@ -1 +0,0 @@
-from .build import build_decoder

xdecoder/body/decoder/build.py

@@ -1,12 +0,0 @@
-from .registry import model_entrypoints
-from .registry import is_model
-
-from .xdecoder import *
-
-def build_decoder(config, *args, **kwargs):
-    model_name = config['MODEL']['DECODER']['NAME']
-
-    if not is_model(model_name):
-        raise ValueError(f'Unkown model: {model_name}')
-
-    return model_entrypoints(model_name)(config, *args, **kwargs)

xdecoder/body/decoder/registry.py

@@ -1,13 +0,0 @@
-_model_entrypoints = {}
-
-def register_decoder(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/body/decoder/tmp.py

@@ -1,664 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-# Modified by Bowen Cheng from: https://github.com/facebookresearch/detr/blob/master/models/detr.py
-import logging
-from typing import Optional
-
-import torch
-from torch import nn, Tensor
-from torch.nn import functional as F
-
-from timm.models.layers import trunc_normal_
-from detectron2.layers import Conv2d
-import fvcore.nn.weight_init as weight_init
-
-from .registry import register_decoder
-from ...utils import configurable
-from ...modules import PositionEmbeddingSine
-
-from image2html.visualizer import VL
-
-
-class SelfAttentionLayer(nn.Module):
-
-    def __init__(self, d_model, nhead, dropout=0.0,
-                 activation="relu", normalize_before=False):
-        super().__init__()
-        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
-
-        self.norm = nn.LayerNorm(d_model)
-        self.dropout = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-        self._reset_parameters()
-    
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(self, tgt,
-                     tgt_mask: Optional[Tensor] = None,
-                     tgt_key_padding_mask: Optional[Tensor] = None,
-                     query_pos: Optional[Tensor] = None):
-        q = k = self.with_pos_embed(tgt, query_pos)
-        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
-                              key_padding_mask=tgt_key_padding_mask)[0]
-        tgt = tgt + self.dropout(tgt2)
-        tgt = self.norm(tgt)
-
-        return tgt
-
-    def forward_pre(self, tgt,
-                    tgt_mask: Optional[Tensor] = None,
-                    tgt_key_padding_mask: Optional[Tensor] = None,
-                    query_pos: Optional[Tensor] = None):
-        tgt2 = self.norm(tgt)
-        q = k = self.with_pos_embed(tgt2, query_pos)
-        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
-                              key_padding_mask=tgt_key_padding_mask)[0]
-        tgt = tgt + self.dropout(tgt2)
-        
-        return tgt
-
-    def forward(self, tgt,
-                tgt_mask: Optional[Tensor] = None,
-                tgt_key_padding_mask: Optional[Tensor] = None,
-                query_pos: Optional[Tensor] = None):
-        if self.normalize_before:
-            return self.forward_pre(tgt, tgt_mask,
-                                    tgt_key_padding_mask, query_pos)
-        return self.forward_post(tgt, tgt_mask,
-                                 tgt_key_padding_mask, query_pos)
-
-
-class CrossAttentionLayer(nn.Module):
-
-    def __init__(self, d_model, nhead, dropout=0.0,
-                 activation="relu", normalize_before=False):
-        super().__init__()
-        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
-
-        self.norm = nn.LayerNorm(d_model)
-        self.dropout = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-        self._reset_parameters()
-    
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(self, tgt, memory,
-                     memory_mask: Optional[Tensor] = None,
-                     memory_key_padding_mask: Optional[Tensor] = None,
-                     pos: Optional[Tensor] = None,
-                     query_pos: Optional[Tensor] = None):
-        tgt2, avg_attn = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
-                                   key=self.with_pos_embed(memory, pos),
-                                   value=memory, attn_mask=memory_mask,
-                                   key_padding_mask=memory_key_padding_mask)
-        tgt = tgt + self.dropout(tgt2)
-        tgt = self.norm(tgt)
-        return tgt, avg_attn
-
-    def forward_pre(self, tgt, memory,
-                    memory_mask: Optional[Tensor] = None,
-                    memory_key_padding_mask: Optional[Tensor] = None,
-                    pos: Optional[Tensor] = None,
-                    query_pos: Optional[Tensor] = None):
-        tgt2 = self.norm(tgt)
-        tgt2, avg_attn = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
-                                   key=self.with_pos_embed(memory, pos),
-                                   value=memory, attn_mask=memory_mask,
-                                   key_padding_mask=memory_key_padding_mask)
-        tgt = tgt + self.dropout(tgt2)
-
-        return tgt, avg_attn
-
-    def forward(self, tgt, memory,
-                memory_mask: Optional[Tensor] = None,
-                memory_key_padding_mask: Optional[Tensor] = None,
-                pos: Optional[Tensor] = None,
-                query_pos: Optional[Tensor] = None):
-        if self.normalize_before:
-            return self.forward_pre(tgt, memory, memory_mask,
-                                    memory_key_padding_mask, pos, query_pos)
-        return self.forward_post(tgt, memory, memory_mask,
-                                 memory_key_padding_mask, pos, query_pos)
-
-
-class FFNLayer(nn.Module):
-
-    def __init__(self, d_model, dim_feedforward=2048, dropout=0.0,
-                 activation="relu", normalize_before=False):
-        super().__init__()
-        # Implementation of Feedforward model
-        self.linear1 = nn.Linear(d_model, dim_feedforward)
-        self.dropout = nn.Dropout(dropout)
-        self.linear2 = nn.Linear(dim_feedforward, d_model)
-
-        self.norm = nn.LayerNorm(d_model)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-        self._reset_parameters()
-    
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(self, tgt):
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
-        tgt = tgt + self.dropout(tgt2)
-        tgt = self.norm(tgt)
-        return tgt
-
-    def forward_pre(self, tgt):
-        tgt2 = self.norm(tgt)
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
-        tgt = tgt + self.dropout(tgt2)
-        return tgt
-
-    def forward(self, tgt):
-        if self.normalize_before:
-            return self.forward_pre(tgt)
-        return self.forward_post(tgt)
-
-
-def _get_activation_fn(activation):
-    """Return an activation function given a string"""
-    if activation == "relu":
-        return F.relu
-    if activation == "gelu":
-        return F.gelu
-    if activation == "glu":
-        return F.glu
-    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
-
-
-class MLP(nn.Module):
-    """ Very simple multi-layer perceptron (also called FFN)"""
-
-    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
-        super().__init__()
-        self.num_layers = num_layers
-        h = [hidden_dim] * (num_layers - 1)
-        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
-
-    def forward(self, x):
-        for i, layer in enumerate(self.layers):
-            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
-        return x
-
-
-class MultiScaleMaskedTransformerDecoder(nn.Module):
-
-    _version = 2
-
-    @configurable
-    def __init__(
-        self,
-        lang_encoder: nn.Module,
-        in_channels,
-        mask_classification=True,
-        *,
-        hidden_dim: int,
-        dim_proj: int,
-        num_queries: int,
-        contxt_len: int,
-        nheads: int,
-        dim_feedforward: int,
-        dec_layers: int,
-        pre_norm: bool,
-        mask_dim: int,
-        task_switch: dict,
-        captioning_step: int,
-        enforce_input_project: bool,
-    ):
-        """
-        NOTE: this interface is experimental.
-        Args:
-            in_channels: channels of the input features
-            mask_classification: whether to add mask classifier or not
-            num_classes: number of classes
-            hidden_dim: Transformer feature dimension
-            num_queries: number of queries
-            nheads: number of heads
-            dim_feedforward: feature dimension in feedforward network
-            enc_layers: number of Transformer encoder layers
-            dec_layers: number of Transformer decoder layers
-            pre_norm: whether to use pre-LayerNorm or not
-            mask_dim: mask feature dimension
-            enforce_input_project: add input project 1x1 conv even if input
-                channels and hidden dim is identical
-        """
-        super().__init__()
-        assert mask_classification, "Only support mask classification model"
-        self.mask_classification = mask_classification
-
-        # positional encoding
-        N_steps = hidden_dim // 2
-        self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
-        
-        # define Transformer decoder here
-        self.num_heads = nheads
-        self.num_layers = dec_layers
-        self.contxt_len = contxt_len
-        self.transformer_self_attention_layers = nn.ModuleList()
-        self.transformer_cross_attention_layers = nn.ModuleList()
-        self.transformer_ffn_layers = nn.ModuleList()
-
-        for _ in range(self.num_layers):
-            self.transformer_self_attention_layers.append(
-                SelfAttentionLayer(
-                    d_model=hidden_dim,
-                    nhead=nheads,
-                    dropout=0.0,
-                    normalize_before=pre_norm,
-                )
-            )
-
-            self.transformer_cross_attention_layers.append(
-                CrossAttentionLayer(
-                    d_model=hidden_dim,
-                    nhead=nheads,
-                    dropout=0.0,
-                    normalize_before=pre_norm,
-                )
-            )
-
-            self.transformer_ffn_layers.append(
-                FFNLayer(
-                    d_model=hidden_dim,
-                    dim_feedforward=dim_feedforward,
-                    dropout=0.0,
-                    normalize_before=pre_norm,
-                )
-            )
-
-        self.decoder_norm = nn.LayerNorm(hidden_dim)
-
-        self.num_queries = num_queries
-        # learnable query features
-        self.query_feat = nn.Embedding(num_queries, hidden_dim)
-        # learnable query p.e.
-        self.query_embed = nn.Embedding(num_queries, hidden_dim)
-        
-        # level embedding (we always use 3 scales)
-        self.num_feature_levels = 3
-        self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim)
-        self.input_proj = nn.ModuleList()
-        
-        for _ in range(self.num_feature_levels):
-            if in_channels != hidden_dim or enforce_input_project:
-                self.input_proj.append(Conv2d(in_channels, hidden_dim, kernel_size=1))
-                weight_init.c2_xavier_fill(self.input_proj[-1])
-            else:
-                self.input_proj.append(nn.Sequential())
-
-        self.task_switch = task_switch
-
-        # output FFNs
-        self.lang_encoder = lang_encoder
-        if self.task_switch['mask']:
-            self.mask_embed = MLP(hidden_dim, hidden_dim, mask_dim, 3)
-
-        self.class_embed = nn.Parameter(torch.empty(hidden_dim, dim_proj))
-        trunc_normal_(self.class_embed, std=.02)
-
-        if task_switch['bbox']:
-            self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
-
-        # Caption Project and query
-        if task_switch['captioning']:
-            self.caping_embed = nn.Parameter(torch.empty(hidden_dim, dim_proj))
-            trunc_normal_(self.caping_embed, std=.02)
-            self.query_feat_caping = nn.Embedding(contxt_len, hidden_dim)
-            self.captioning_step = captioning_step
-
-        # register self_attn_mask to avoid information leakage, it includes interaction between object query, class query and caping query
-        self_attn_mask = torch.zeros((1, num_queries + contxt_len, num_queries + contxt_len)).bool()
-        self_attn_mask[:, :num_queries, num_queries:] = True # object+class query does not attend with caption query.
-        self_attn_mask[:, num_queries:, num_queries:] = torch.triu(torch.ones((1, contxt_len, contxt_len)), diagonal=1).bool() # caption query only attend with previous token.
-        self_attn_mask[:, :num_queries-1, num_queries-1:num_queries] = True # object query does not attend with class query.
-        self_attn_mask[:, num_queries-1:num_queries, :num_queries-1] = True # class query does not attend with object query.
-        self.register_buffer("self_attn_mask", self_attn_mask)
-
-
-    @classmethod
-    def from_config(cls, cfg, in_channels, lang_encoder, mask_classification, extra):
-        ret = {}
-
-        ret["lang_encoder"] = lang_encoder
-        ret["in_channels"] = in_channels
-        ret["mask_classification"] = mask_classification
-
-        enc_cfg = cfg['MODEL']['ENCODER']
-        dec_cfg = cfg['MODEL']['DECODER']
-        
-        ret["hidden_dim"] = dec_cfg['HIDDEN_DIM']
-        ret["dim_proj"] = cfg['MODEL']['DIM_PROJ']
-        ret["num_queries"] = dec_cfg['NUM_OBJECT_QUERIES']
-        ret["contxt_len"] = cfg['MODEL']['TEXT']['CONTEXT_LENGTH']
-        
-        # Transformer parameters:
-        ret["nheads"] = dec_cfg['NHEADS']
-        ret["dim_feedforward"] = dec_cfg['DIM_FEEDFORWARD']
-
-        # NOTE: because we add learnable query features which requires supervision,
-        # we add minus 1 to decoder layers to be consistent with our loss
-        # implementation: that is, number of auxiliary losses is always
-        # equal to number of decoder layers. With learnable query features, the number of
-        # auxiliary losses equals number of decoders plus 1.
-        assert dec_cfg['DEC_LAYERS'] >= 1
-        ret["dec_layers"] = dec_cfg['DEC_LAYERS'] - 1
-        ret["pre_norm"] = dec_cfg['PRE_NORM']
-        ret["enforce_input_project"] = dec_cfg['ENFORCE_INPUT_PROJ']
-        ret["mask_dim"] = enc_cfg['MASK_DIM']
-
-        ret["task_switch"] = extra['task_switch']
-        ret["captioning_step"] = dec_cfg['CAPTIONING'].get('STEP', 50)
-
-        return ret
-
-    def forward(self, x, mask_features, mask=None, target_queries=None, target_vlp=None, task='seg', extra={}):
-        if task == 'captioning_infer':
-            return self.forward_captioning(x, mask_features, mask=mask, target_queries=target_queries, target_vlp=target_vlp, task=task, extra=extra)
-        # x is a list of multi-scale feature
-        assert len(x) == self.num_feature_levels
-        src = []
-        pos = []
-        size_list = []
-        
-        # disable mask, it does not affect performance
-        del mask
-        for i in range(self.num_feature_levels):
-            size_list.append(x[i].shape[-2:])
-            pos.append(self.pe_layer(x[i], None).flatten(2))
-            src.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
-
-            # flatten NxCxHxW to HWxNxC
-            pos[-1] = pos[-1].permute(2, 0, 1)
-            src[-1] = src[-1].permute(2, 0, 1)
-
-        _, bs, _ = src[0].shape
-
-        # QxNxC
-        query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
-        output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)
-
-        predictions_class = []
-        predictions_mask = []
-        predictions_bbox = []
-        predictions_caption = []
-        predictions_captioning = []
-        
-        self_tgt_mask = None
-        if self.training and task == 'vlp' and self.task_switch['captioning']:
-            output = torch.cat((output, self.query_feat_caping.weight.unsqueeze(1).repeat(1, bs, 1)), dim=0) # concat object query, class token and caption token.
-            caping_lang_embed = torch.cat([caption['caption_tokens'] for caption in target_vlp], dim=0).transpose(0, 1) # language output
-            query_embed = torch.cat((query_embed, caping_lang_embed), dim=0) # may not add at the beginning.
-            self_tgt_mask = self.self_attn_mask.repeat(output.shape[1]*self.num_heads, 1, 1)
-        elif (((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding']) \
-                or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-            self_tgt_mask = self.self_attn_mask[:,:self.num_queries,:self.num_queries].repeat(output.shape[1]*self.num_heads, 1, 1)
-            grounding_tokens = extra['grounding_tokens']
-            _grounding_tokens = grounding_tokens.detach().clone()
-            # initialize with negative attention at the beginning.
-            pad_tgt_mask = torch.ones((1, self.num_queries + (self.num_queries-1) + len(grounding_tokens), self.num_queries + (self.num_queries-1) + len(grounding_tokens)), device=self_tgt_mask.device).bool().repeat(output.shape[1]*self.num_heads, 1, 1)
-            pad_tgt_mask[:,:self.num_queries,:self.num_queries] = self_tgt_mask
-            pad_tgt_mask[:,self.num_queries:,self.num_queries:] = False # grounding tokens could attend with eatch other
-            self_tgt_mask = pad_tgt_mask
-            output = torch.cat((output, output[:-1]), dim=0)
-            query_embed = torch.cat((query_embed, query_embed[:-1]), dim=0) # also pad language embdding to fix embedding
-        else:
-            self_tgt_mask = self.self_attn_mask[:,:self.num_queries,:self.num_queries].repeat(output.shape[1]*self.num_heads, 1, 1)
-
-        # prediction heads on learnable query features
-        results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[0], task=task)
-        attn_mask = results["attn_mask"]
-        predictions_class.append(results["outputs_class"])
-        predictions_mask.append(results["outputs_mask"])
-        predictions_bbox.append(results["outputs_bbox"])
-        predictions_caption.append(results["outputs_caption"])
-        predictions_captioning.append(results["outputs_captionting"])
-        
-        for i in range(self.num_layers):
-            level_index = i % self.num_feature_levels
-            attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
-
-            if self.training and task == 'vlp' and self.task_switch['captioning']:
-                attn_mask = torch.cat((attn_mask, torch.zeros_like(attn_mask[:, :self.contxt_len, :])), dim=1)
-            # attention: cross-attention first
-            output, avg_attn = self.transformer_cross_attention_layers[i](
-                output, src[level_index],
-                memory_mask=attn_mask,
-                memory_key_padding_mask=None,  # here we do not apply masking on padded region
-                pos=pos[level_index], query_pos=query_embed
-            )
-
-            if (((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding']) \
-                    or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-                output = torch.cat((output, _grounding_tokens), dim=0)
-                query_embed = torch.cat((query_embed, grounding_tokens), dim=0)
-
-            output = self.transformer_self_attention_layers[i](
-                output, tgt_mask=self_tgt_mask,
-                tgt_key_padding_mask=None,
-                query_pos=query_embed
-            )
-            
-            # FFN
-            output = self.transformer_ffn_layers[i](
-                output
-            )
-
-            if ((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding'] \
-                    or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-                _grounding_tokens = output[-len(_grounding_tokens):]
-                output = output[:-len(_grounding_tokens)]
-                query_embed = query_embed[:-len(_grounding_tokens)]
-
-            results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels], layer_id=i, task=task)
-            attn_mask = results["attn_mask"]
-            predictions_class.append(results["outputs_class"])
-            predictions_mask.append(results["outputs_mask"])
-            predictions_bbox.append(results["outputs_bbox"])
-            predictions_caption.append(results["outputs_caption"])
-            predictions_captioning.append(results["outputs_captionting"])
-
-        assert len(predictions_class) == self.num_layers + 1
-        if task == 'vlp':
-            out = {'pred_captionings': predictions_captioning[-1], 
-                   'pred_captions': predictions_caption[-1], 
-                   'aux_outputs': [{'pred_captionings': x, 'pred_captions': y } for x, y in zip(predictions_captioning[:-1], predictions_caption[:-1])]}
-            return out
-        else:
-            out = {
-                'pred_logits': predictions_class[-1],
-                'pred_masks': predictions_mask[-1],
-                'pred_boxes': predictions_bbox[-1],
-                'pred_captions': predictions_caption[-1],
-                'aux_outputs': self._set_aux_loss(
-                    predictions_class if self.mask_classification else None, predictions_mask, predictions_bbox, predictions_caption
-                )
-            }
-            return out
-
-    def forward_captioning(self, x, mask_features, mask = None, target_queries = None, target_vlp = None, task='seg', extra={}):
-        # x is a list of multi-scale feature
-        assert len(x) == self.num_feature_levels
-        src = []
-        pos = []
-        size_list = []
-        
-        # disable mask, it does not affect performance
-        del mask
-        for i in range(self.num_feature_levels):
-            size_list.append(x[i].shape[-2:])
-            pos.append(self.pe_layer(x[i], None).flatten(2))
-            src.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
-
-            # flatten NxCxHxW to HWxNxC
-            pos[-1] = pos[-1].permute(2, 0, 1)
-            src[-1] = src[-1].permute(2, 0, 1)
-
-        _, bs, _ = src[0].shape
-
-        # QxNxC
-        query_embed_ = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
-        query_feat = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)        
-        caping_lang_token = extra['start_token'].repeat(bs, 1)
-        query_feat_caping = self.query_feat_caping.weight.unsqueeze(1).repeat(1, bs, 1)
-        
-        # prepare token embedding for evaluation
-        token_embs = self.lang_encoder.lang_encoder.token_embedding.weight
-        # token_embs = (token_embs / token_embs.norm(dim=-1, keepdim=True) + 1e-7)
-        
-        for cap_idx in range(0, self.captioning_step):
-            caping_lang_embed = self.lang_encoder.forward_language_token((caping_lang_token,))[0].transpose(0, 1)
-            query_embed = torch.cat((query_embed_, caping_lang_embed), dim=0) # may not add at the beginning.
-            output = torch.cat((query_feat, query_feat_caping), dim=0) # concat object query, class token and caption token.
-
-            # prediction heads on learnable query features
-            results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[0], task=task)
-            attn_mask = results["attn_mask"]
-        
-            for i in range(self.num_layers):
-                level_index = i % self.num_feature_levels
-                attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
-                attn_mask = torch.cat((attn_mask, torch.zeros_like(attn_mask[:, :self.contxt_len, :])), dim=1)
-                self_tgt_mask = self.self_attn_mask.repeat(output.shape[1]*self.num_heads, 1, 1)
-                
-                # attention: cross-attention first
-                output, avg_attn = self.transformer_cross_attention_layers[i](
-                    output, src[level_index],
-                    memory_mask=attn_mask,
-                    memory_key_padding_mask=None,  # here we do not apply masking on padded region
-                    pos=pos[level_index], query_pos=query_embed
-                )
-
-                output = self.transformer_self_attention_layers[i](
-                    output, tgt_mask=self_tgt_mask,
-                    tgt_key_padding_mask=None,
-                    query_pos=query_embed
-                )
-                
-                # FFN
-                output = self.transformer_ffn_layers[i](
-                    output
-                )
-
-                results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels], layer_id=i, task=task)
-                attn_mask = results["attn_mask"]
-            
-            pred_captions_gen = results['outputs_captionting']
-            # pred_captions_gen = (pred_captions_gen / pred_captions_gen.norm(dim=-1, keepdim=True) + 1e-7)
-            pred_captions_gen = pred_captions_gen @ token_embs.t()
-            caping_lang_token[:,cap_idx+1] = pred_captions_gen[:,cap_idx].max(-1)[1]
-
-        out = {'pred_captionings': caping_lang_token,
-               'pred_texts': self.lang_encoder.tokenizer.batch_decode(caping_lang_token, skip_special_tokens=True)}
-        return out
-
-
-    def forward_prediction_heads(self, output, mask_features, attn_mask_target_size, layer_id=-1, task='seg'):
-        decoder_output = self.decoder_norm(output)
-        decoder_output = decoder_output.transpose(0, 1)
-
-        # extract image captioning token from decoder output.
-        if self.task_switch['captioning'] and (task == 'vlp' or task == 'captioning_infer'):
-            outputs_captionting = decoder_output[:,self.num_queries:] @ self.caping_embed
-        else:
-            outputs_captionting = None
-
-        # recompute class token output.
-        norm_decoder_output = decoder_output / (decoder_output.norm(dim=-1, keepdim=True) + 1e-7)
-        obj_token = norm_decoder_output[:,:self.num_queries-1]
-        cls_token = norm_decoder_output[:,self.num_queries-1:self.num_queries]
-
-        sim = (cls_token @ obj_token.transpose(1,2)).softmax(-1)[:,0,:,None] # TODO include class token.
-        cls_token = (sim * decoder_output[:,:self.num_queries-1]).sum(dim=1, keepdim=True)
-
-        if (((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding']) \
-                or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-            decoder_output = torch.cat((decoder_output[:,:self.num_queries-1], cls_token, decoder_output[:,self.num_queries:2*self.num_queries-1]), dim=1)
-        else:
-            decoder_output = torch.cat((decoder_output[:,:self.num_queries-1], cls_token), dim=1)
-
-        # compute class, mask and bbox.
-        class_embed = decoder_output @ self.class_embed
-        # HACK do not compute similarity if mask is not on
-        outputs_class = self.lang_encoder.compute_similarity(class_embed, fake=(((not self.task_switch['mask']) and self.training) or (task == 'openimage')))
-
-        if self.task_switch['mask'] or self.task_switch['openimage']['mask']:
-            mask_embed = self.mask_embed(decoder_output)
-            outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)
-
-            # NOTE: prediction is of higher-resolution
-            # [B, Q, H, W] -> [B, Q, H*W] -> [B, h, Q, H*W] -> [B*h, Q, HW]
-            attn_mask = F.interpolate(outputs_mask, size=attn_mask_target_size, mode="bilinear", align_corners=False)
-
-            # must use bool type
-            # If a BoolTensor is provided, positions with ``True`` are not allowed to attend while ``False`` values will be unchanged.
-            attn_mask = (attn_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5).bool()
-            attn_mask = attn_mask.detach()
-
-            # NOTE: fill False for cls token (JY)
-            attn_mask[:, self.num_queries:self.num_queries+1].fill_(False)
-        else:
-            outputs_mask = None
-            attn_mask = torch.zeros((list(decoder_output.shape[:2]) + [attn_mask_target_size[0]*attn_mask_target_size[1]]), device=decoder_output.device).repeat(self.num_heads, 1, 1).bool()
-
-        outputs_bbox = [None for i in range(len(decoder_output))]
-        if self.task_switch['bbox']:
-            outputs_bbox = self.bbox_embed(decoder_output)
-
-        outputs_caption = None
-        if self.task_switch['caption']:
-            outputs_caption = class_embed
-            
-
-        results = {
-            "outputs_class": outputs_class,
-            "outputs_mask": outputs_mask,
-            "outputs_bbox": outputs_bbox,
-            "attn_mask": attn_mask,
-            "outputs_caption": outputs_caption,
-            "outputs_captionting": outputs_captionting,
-        }
-        return results
-
-    @torch.jit.unused
-    def _set_aux_loss(self, outputs_class, outputs_seg_masks, outputs_boxes, outputs_captions):
-        # this is a workaround to make torchscript happy, as torchscript
-        # doesn't support dictionary with non-homogeneous values, such
-        # as a dict having both a Tensor and a list.
-        if self.mask_classification:
-            return [
-                {"pred_logits": a, "pred_masks": b, "pred_boxes": c, "pred_captions": d}
-                for a, b, c, d in zip(outputs_class[:-1], outputs_seg_masks[:-1], outputs_boxes[:-1], outputs_captions[:-1])
-            ]
-        else:
-            return [{"pred_masks": b} for b in outputs_seg_masks[:-1]]
-
-
-@register_decoder
-def get_masked_transformer_decoder(cfg, in_channels, lang_encoder, mask_classification, extra):
-    return MultiScaleMaskedTransformerDecoder(cfg, in_channels, lang_encoder, mask_classification, extra)

xdecoder/body/decoder/xdecoder.py

@@ -1,700 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-# Modified by Bowen Cheng from: https://github.com/facebookresearch/detr/blob/master/models/detr.py
-
-# --------------------------------------------------------
-# X-Decoder -- Generalized Decoding for Pixel, Image, and Language
-# Copyright (c) 2022 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Xueyan Zou (xueyan@cs.wisc.edu), Jianwei Yang (jianwyan@microsoft.com)
-# --------------------------------------------------------
-
-
-import logging
-from typing import Optional
-
-import torch
-from torch import nn, Tensor
-from torch.nn import functional as F
-
-from timm.models.layers import trunc_normal_
-from detectron2.layers import Conv2d
-import fvcore.nn.weight_init as weight_init
-
-from .registry import register_decoder
-from ...utils import configurable
-from ...modules import PositionEmbeddingSine
-
-
-class SelfAttentionLayer(nn.Module):
-
-    def __init__(self, d_model, nhead, dropout=0.0,
-                 activation="relu", normalize_before=False):
-        super().__init__()
-        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
-
-        self.norm = nn.LayerNorm(d_model)
-        self.dropout = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-        self._reset_parameters()
-    
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(self, tgt,
-                     tgt_mask: Optional[Tensor] = None,
-                     tgt_key_padding_mask: Optional[Tensor] = None,
-                     query_pos: Optional[Tensor] = None):
-        q = k = self.with_pos_embed(tgt, query_pos)
-        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
-                              key_padding_mask=tgt_key_padding_mask)[0]
-        tgt = tgt + self.dropout(tgt2)
-        tgt = self.norm(tgt)
-
-        return tgt
-
-    def forward_pre(self, tgt,
-                    tgt_mask: Optional[Tensor] = None,
-                    tgt_key_padding_mask: Optional[Tensor] = None,
-                    query_pos: Optional[Tensor] = None):
-        tgt2 = self.norm(tgt)
-        q = k = self.with_pos_embed(tgt2, query_pos)
-        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
-                              key_padding_mask=tgt_key_padding_mask)[0]
-        tgt = tgt + self.dropout(tgt2)
-        
-        return tgt
-
-    def forward(self, tgt,
-                tgt_mask: Optional[Tensor] = None,
-                tgt_key_padding_mask: Optional[Tensor] = None,
-                query_pos: Optional[Tensor] = None):
-        if self.normalize_before:
-            return self.forward_pre(tgt, tgt_mask,
-                                    tgt_key_padding_mask, query_pos)
-        return self.forward_post(tgt, tgt_mask,
-                                 tgt_key_padding_mask, query_pos)
-
-
-class CrossAttentionLayer(nn.Module):
-
-    def __init__(self, d_model, nhead, dropout=0.0,
-                 activation="relu", normalize_before=False):
-        super().__init__()
-        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
-
-        self.norm = nn.LayerNorm(d_model)
-        self.dropout = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-        self._reset_parameters()
-    
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(self, tgt, memory,
-                     memory_mask: Optional[Tensor] = None,
-                     memory_key_padding_mask: Optional[Tensor] = None,
-                     pos: Optional[Tensor] = None,
-                     query_pos: Optional[Tensor] = None):
-        tgt2, avg_attn = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
-                                   key=self.with_pos_embed(memory, pos),
-                                   value=memory, attn_mask=memory_mask,
-                                   key_padding_mask=memory_key_padding_mask)
-        tgt = tgt + self.dropout(tgt2)
-        tgt = self.norm(tgt)
-        return tgt, avg_attn
-
-    def forward_pre(self, tgt, memory,
-                    memory_mask: Optional[Tensor] = None,
-                    memory_key_padding_mask: Optional[Tensor] = None,
-                    pos: Optional[Tensor] = None,
-                    query_pos: Optional[Tensor] = None):
-        tgt2 = self.norm(tgt)
-        tgt2, avg_attn = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
-                                   key=self.with_pos_embed(memory, pos),
-                                   value=memory, attn_mask=memory_mask,
-                                   key_padding_mask=memory_key_padding_mask)
-        tgt = tgt + self.dropout(tgt2)
-
-        return tgt, avg_attn
-
-    def forward(self, tgt, memory,
-                memory_mask: Optional[Tensor] = None,
-                memory_key_padding_mask: Optional[Tensor] = None,
-                pos: Optional[Tensor] = None,
-                query_pos: Optional[Tensor] = None):
-        if self.normalize_before:
-            return self.forward_pre(tgt, memory, memory_mask,
-                                    memory_key_padding_mask, pos, query_pos)
-        return self.forward_post(tgt, memory, memory_mask,
-                                 memory_key_padding_mask, pos, query_pos)
-
-
-class FFNLayer(nn.Module):
-
-    def __init__(self, d_model, dim_feedforward=2048, dropout=0.0,
-                 activation="relu", normalize_before=False):
-        super().__init__()
-        # Implementation of Feedforward model
-        self.linear1 = nn.Linear(d_model, dim_feedforward)
-        self.dropout = nn.Dropout(dropout)
-        self.linear2 = nn.Linear(dim_feedforward, d_model)
-
-        self.norm = nn.LayerNorm(d_model)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-        self._reset_parameters()
-    
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(self, tgt):
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
-        tgt = tgt + self.dropout(tgt2)
-        tgt = self.norm(tgt)
-        return tgt
-
-    def forward_pre(self, tgt):
-        tgt2 = self.norm(tgt)
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
-        tgt = tgt + self.dropout(tgt2)
-        return tgt
-
-    def forward(self, tgt):
-        if self.normalize_before:
-            return self.forward_pre(tgt)
-        return self.forward_post(tgt)
-
-
-def _get_activation_fn(activation):
-    """Return an activation function given a string"""
-    if activation == "relu":
-        return F.relu
-    if activation == "gelu":
-        return F.gelu
-    if activation == "glu":
-        return F.glu
-    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
-
-
-class MLP(nn.Module):
-    """ Very simple multi-layer perceptron (also called FFN)"""
-
-    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
-        super().__init__()
-        self.num_layers = num_layers
-        h = [hidden_dim] * (num_layers - 1)
-        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
-
-    def forward(self, x):
-        for i, layer in enumerate(self.layers):
-            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
-        return x
-
-
-class MultiScaleMaskedTransformerDecoder(nn.Module):
-
-    _version = 2
-
-    @configurable
-    def __init__(
-        self,
-        lang_encoder: nn.Module,
-        in_channels,
-        mask_classification=True,
-        *,
-        hidden_dim: int,
-        dim_proj: int,
-        num_queries: int,
-        contxt_len: int,
-        nheads: int,
-        dim_feedforward: int,
-        dec_layers: int,
-        pre_norm: bool,
-        mask_dim: int,
-        task_switch: dict,
-        captioning_step: int,
-        enforce_input_project: bool,
-    ):
-        """
-        NOTE: this interface is experimental.
-        Args:
-            in_channels: channels of the input features
-            mask_classification: whether to add mask classifier or not
-            num_classes: number of classes
-            hidden_dim: Transformer feature dimension
-            num_queries: number of queries
-            nheads: number of heads
-            dim_feedforward: feature dimension in feedforward network
-            enc_layers: number of Transformer encoder layers
-            dec_layers: number of Transformer decoder layers
-            pre_norm: whether to use pre-LayerNorm or not
-            mask_dim: mask feature dimension
-            enforce_input_project: add input project 1x1 conv even if input
-                channels and hidden dim is identical
-        """
-        super().__init__()
-        assert mask_classification, "Only support mask classification model"
-        self.mask_classification = mask_classification
-
-        # positional encoding
-        N_steps = hidden_dim // 2
-        self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
-        
-        # define Transformer decoder here
-        self.num_heads = nheads
-        self.num_layers = dec_layers
-        self.contxt_len = contxt_len
-        self.transformer_self_attention_layers = nn.ModuleList()
-        self.transformer_cross_attention_layers = nn.ModuleList()
-        self.transformer_ffn_layers = nn.ModuleList()
-
-        for _ in range(self.num_layers):
-            self.transformer_self_attention_layers.append(
-                SelfAttentionLayer(
-                    d_model=hidden_dim,
-                    nhead=nheads,
-                    dropout=0.0,
-                    normalize_before=pre_norm,
-                )
-            )
-
-            self.transformer_cross_attention_layers.append(
-                CrossAttentionLayer(
-                    d_model=hidden_dim,
-                    nhead=nheads,
-                    dropout=0.0,
-                    normalize_before=pre_norm,
-                )
-            )
-
-            self.transformer_ffn_layers.append(
-                FFNLayer(
-                    d_model=hidden_dim,
-                    dim_feedforward=dim_feedforward,
-                    dropout=0.0,
-                    normalize_before=pre_norm,
-                )
-            )
-
-        self.decoder_norm = nn.LayerNorm(hidden_dim)
-
-        self.num_queries = num_queries
-        # learnable query features
-        self.query_feat = nn.Embedding(num_queries, hidden_dim)
-        # learnable query p.e.
-        self.query_embed = nn.Embedding(num_queries, hidden_dim)
-        
-        # level embedding (we always use 3 scales)
-        self.num_feature_levels = 3
-        self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim)
-        self.input_proj = nn.ModuleList()
-        
-        for _ in range(self.num_feature_levels):
-            if in_channels != hidden_dim or enforce_input_project:
-                self.input_proj.append(Conv2d(in_channels, hidden_dim, kernel_size=1))
-                weight_init.c2_xavier_fill(self.input_proj[-1])
-            else:
-                self.input_proj.append(nn.Sequential())
-
-        self.task_switch = task_switch
-
-        # output FFNs
-        self.lang_encoder = lang_encoder
-        if self.task_switch['mask']:
-            self.mask_embed = MLP(hidden_dim, hidden_dim, mask_dim, 3)
-
-        self.class_embed = nn.Parameter(torch.empty(hidden_dim, dim_proj))
-        trunc_normal_(self.class_embed, std=.02)
-
-        if task_switch['bbox']:
-            self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
-
-        # Caption Project and query
-        if task_switch['captioning']:
-            self.caping_embed = nn.Parameter(torch.empty(hidden_dim, dim_proj))
-            trunc_normal_(self.caping_embed, std=.02)
-            # self.query_feat_caping = nn.Embedding(contxt_len, hidden_dim)
-            self.pos_embed_caping = nn.Embedding(contxt_len, hidden_dim)
-            self.captioning_step = captioning_step
-
-        # register self_attn_mask to avoid information leakage, it includes interaction between object query, class query and caping query
-        self_attn_mask = torch.zeros((1, num_queries + contxt_len, num_queries + contxt_len)).bool()
-        self_attn_mask[:, :num_queries, num_queries:] = True # object+class query does not attend with caption query.
-        self_attn_mask[:, num_queries:, num_queries:] = torch.triu(torch.ones((1, contxt_len, contxt_len)), diagonal=1).bool() # caption query only attend with previous token.
-        self_attn_mask[:, :num_queries-1, num_queries-1:num_queries] = True # object query does not attend with class query.
-        self_attn_mask[:, num_queries-1:num_queries, :num_queries-1] = True # class query does not attend with object query.
-        self.register_buffer("self_attn_mask", self_attn_mask)
-
-
-    @classmethod
-    def from_config(cls, cfg, in_channels, lang_encoder, mask_classification, extra):
-        ret = {}
-
-        ret["lang_encoder"] = lang_encoder
-        ret["in_channels"] = in_channels
-        ret["mask_classification"] = mask_classification
-
-        enc_cfg = cfg['MODEL']['ENCODER']
-        dec_cfg = cfg['MODEL']['DECODER']
-        
-        ret["hidden_dim"] = dec_cfg['HIDDEN_DIM']
-        ret["dim_proj"] = cfg['MODEL']['DIM_PROJ']
-        ret["num_queries"] = dec_cfg['NUM_OBJECT_QUERIES']
-        ret["contxt_len"] = cfg['MODEL']['TEXT']['CONTEXT_LENGTH']
-        
-        # Transformer parameters:
-        ret["nheads"] = dec_cfg['NHEADS']
-        ret["dim_feedforward"] = dec_cfg['DIM_FEEDFORWARD']
-
-        # NOTE: because we add learnable query features which requires supervision,
-        # we add minus 1 to decoder layers to be consistent with our loss
-        # implementation: that is, number of auxiliary losses is always
-        # equal to number of decoder layers. With learnable query features, the number of
-        # auxiliary losses equals number of decoders plus 1.
-        assert dec_cfg['DEC_LAYERS'] >= 1
-        ret["dec_layers"] = dec_cfg['DEC_LAYERS'] - 1
-        ret["pre_norm"] = dec_cfg['PRE_NORM']
-        ret["enforce_input_project"] = dec_cfg['ENFORCE_INPUT_PROJ']
-        ret["mask_dim"] = enc_cfg['MASK_DIM']
-
-        ret["task_switch"] = extra['task_switch']
-        ret["captioning_step"] = dec_cfg['CAPTIONING'].get('STEP', 50)
-
-        return ret
-
-    def forward(self, x, mask_features, mask=None, target_queries=None, target_vlp=None, task='seg', extra={}):
-        if task == 'captioning_infer':
-            return self.forward_captioning(x, mask_features, mask=mask, target_queries=target_queries, target_vlp=target_vlp, task=task, extra=extra)
-        # x is a list of multi-scale feature
-        assert len(x) == self.num_feature_levels
-        src = []
-        pos = []
-        size_list = []
-        
-        # disable mask, it does not affect performance
-        del mask
-        for i in range(self.num_feature_levels):
-            size_list.append(x[i].shape[-2:])
-            pos.append(self.pe_layer(x[i], None).flatten(2))
-            src.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
-
-            # flatten NxCxHxW to HWxNxC
-            pos[-1] = pos[-1].permute(2, 0, 1)
-            src[-1] = src[-1].permute(2, 0, 1)
-
-        _, bs, _ = src[0].shape
-
-        # QxNxC
-        query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
-        output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)
-
-        predictions_class = []
-        predictions_mask = []
-        predictions_bbox = []
-        predictions_caption = []
-        predictions_captioning = []
-        
-        self_tgt_mask = None
-        if self.training and task == 'vlp' and self.task_switch['captioning']:
-            # output = torch.cat((output, self.query_feat_caping.weight.unsqueeze(1).repeat(1, bs, 1)), dim=0) # concat object query, class token and caption token.
-            caping_lang_embed = torch.cat([caption['caption_tokens'] for caption in target_vlp], dim=0).transpose(0, 1) # language output
-            _caping_lang_embed = caping_lang_embed.detach().clone()
-            output = torch.cat((output, _caping_lang_embed), dim=0) # concat object query, class token and caption token.
-            caping_lang_embed += self.pos_embed_caping.weight.unsqueeze(1).repeat(1, bs, 1)
-            query_embed = torch.cat((query_embed, caping_lang_embed), dim=0) # may not add at the beginning.
-            self_tgt_mask = self.self_attn_mask.repeat(output.shape[1]*self.num_heads, 1, 1)
-        elif (((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding']) \
-                or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-            self_tgt_mask = self.self_attn_mask[:,:self.num_queries,:self.num_queries].repeat(output.shape[1]*self.num_heads, 1, 1)
-            grounding_tokens = extra['grounding_tokens']
-            _grounding_tokens = grounding_tokens.detach().clone()
-            # initialize with negative attention at the beginning.
-            pad_tgt_mask = torch.ones((1, self.num_queries + (self.num_queries-1) + len(grounding_tokens), self.num_queries + (self.num_queries-1) + len(grounding_tokens)), device=self_tgt_mask.device).bool().repeat(output.shape[1]*self.num_heads, 1, 1)
-            pad_tgt_mask[:,:self.num_queries,:self.num_queries] = self_tgt_mask
-            pad_tgt_mask[:,self.num_queries:,self.num_queries:] = False # grounding tokens could attend with eatch other
-            self_tgt_mask = pad_tgt_mask
-            output = torch.cat((output, output[:-1]), dim=0)
-            query_embed = torch.cat((query_embed, query_embed[:-1]), dim=0) # also pad language embdding to fix embedding
-        else:
-            self_tgt_mask = self.self_attn_mask[:,:self.num_queries,:self.num_queries].repeat(output.shape[1]*self.num_heads, 1, 1)
-
-        # prediction heads on learnable query features
-        results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[0], task=task)
-        attn_mask = results["attn_mask"]
-        predictions_class.append(results["outputs_class"])
-        predictions_mask.append(results["outputs_mask"])
-        predictions_bbox.append(results["outputs_bbox"])
-        predictions_caption.append(results["outputs_caption"])
-        predictions_captioning.append(results["outputs_captionting"])
-        
-        for i in range(self.num_layers):
-            level_index = i % self.num_feature_levels
-            attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
-
-            if self.training and task == 'vlp' and self.task_switch['captioning']:
-                attn_mask = torch.cat((attn_mask, torch.zeros_like(attn_mask[:, :self.contxt_len, :])), dim=1)
-            # attention: cross-attention first
-            output, avg_attn = self.transformer_cross_attention_layers[i](
-                output, src[level_index],
-                memory_mask=attn_mask,
-                memory_key_padding_mask=None,  # here we do not apply masking on padded region
-                pos=pos[level_index], query_pos=query_embed
-            )
-
-            if (((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding']) \
-                    or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-                output = torch.cat((output, _grounding_tokens), dim=0)
-                query_embed = torch.cat((query_embed, grounding_tokens), dim=0)
-
-            output = self.transformer_self_attention_layers[i](
-                output, tgt_mask=self_tgt_mask,
-                tgt_key_padding_mask=None,
-                query_pos=query_embed
-            )
-            
-            # FFN
-            output = self.transformer_ffn_layers[i](
-                output
-            )
-
-            if ((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding'] \
-                    or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-                _grounding_tokens = output[-len(_grounding_tokens):]
-                output = output[:-len(_grounding_tokens)]
-                query_embed = query_embed[:-len(_grounding_tokens)]
-
-            results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels], layer_id=i, task=task)
-            attn_mask = results["attn_mask"]
-            predictions_class.append(results["outputs_class"])
-            predictions_mask.append(results["outputs_mask"])
-            predictions_bbox.append(results["outputs_bbox"])
-            predictions_caption.append(results["outputs_caption"])
-            predictions_captioning.append(results["outputs_captionting"])
-
-        assert len(predictions_class) == self.num_layers + 1
-        if task == 'vlp':
-            out = {'pred_captionings': predictions_captioning[-1], 
-                   'pred_captions': predictions_caption[-1], 
-                   'aux_outputs': [{'pred_captionings': x, 'pred_captions': y } for x, y in zip(predictions_captioning[:-1], predictions_caption[:-1])]}
-            return out
-        else:
-            out = {
-                'pred_logits': predictions_class[-1],
-                'pred_masks': predictions_mask[-1],
-                'pred_boxes': predictions_bbox[-1],
-                'pred_captions': predictions_caption[-1],
-                'aux_outputs': self._set_aux_loss(
-                    predictions_class if self.mask_classification else None, predictions_mask, predictions_bbox, predictions_caption
-                )
-            }
-            return out
-
-    def forward_captioning(self, x, mask_features, mask = None, target_queries = None, target_vlp = None, task='seg', extra={}):
-        # x is a list of multi-scale feature
-        assert len(x) == self.num_feature_levels
-        src = []
-        pos = []
-        size_list = []
-        
-        # disable mask, it does not affect performance
-        del mask
-        for i in range(self.num_feature_levels):
-            size_list.append(x[i].shape[-2:])
-            pos.append(self.pe_layer(x[i], None).flatten(2))
-            src.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
-
-            # flatten NxCxHxW to HWxNxC
-            pos[-1] = pos[-1].permute(2, 0, 1)
-            src[-1] = src[-1].permute(2, 0, 1)
-
-        _, bs, _ = src[0].shape
-
-        # QxNxC
-        query_embed_ = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
-        query_feat = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)        
-        caping_lang_token = extra['start_token'].repeat(bs, 1)
-        start_id = 0
-        if 'token' in extra:
-            caping_lang_token[:,:len(extra['token'][0])] = extra['token']
-            start_id = len(extra['token'][0])-1
-        # query_feat_caping = self.query_feat_caping.weight.unsqueeze(1).repeat(1, bs, 1)
-        pos_embed_caping = self.pos_embed_caping.weight.unsqueeze(1).repeat(1, bs, 1)
-        # prepare token embedding for evaluation
-        token_embs = self.lang_encoder.lang_encoder.token_embedding.weight
-        # token_embs = (token_embs / token_embs.norm(dim=-1, keepdim=True) + 1e-7)
-        
-        for cap_idx in range(start_id, self.captioning_step):
-            caping_lang_embed = self.lang_encoder.forward_language_token((caping_lang_token,))[0].transpose(0, 1)
-            output = torch.cat((query_feat, caping_lang_embed), dim=0) # concat object query, class token and caption token.
-            caping_lang_embed += pos_embed_caping
-            query_embed = torch.cat((query_embed_, caping_lang_embed), dim=0) # may not add at the beginning.
-            # output = torch.cat((query_feat, query_feat_caping), dim=0) # concat object query, class token and caption token.
-
-            # prediction heads on learnable query features
-            results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[0], task=task)
-            attn_mask = results["attn_mask"]
-        
-            for i in range(self.num_layers):
-                level_index = i % self.num_feature_levels
-                attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
-                attn_mask = torch.cat((attn_mask, torch.zeros_like(attn_mask[:, :self.contxt_len, :])), dim=1)
-                self_tgt_mask = self.self_attn_mask.repeat(output.shape[1]*self.num_heads, 1, 1)
-
-                if extra['captioning_mask'] is not None:
-                    bs,nq,wh = attn_mask.shape
-                    assert bs==self.num_heads, "Only support single image referring captioning."
-                    cap_mask = extra['captioning_mask']
-                    attn_mask = attn_mask.reshape(bs,nq,size_list[i%3][0],size_list[i%3][1])
-                    cap_mask = F.interpolate(cap_mask[None,].float(), size_list[i%3], mode='nearest').bool()[0,0]
-                    attn_mask[:,self.num_queries:, cap_mask] = True
-                    attn_mask = attn_mask.reshape(bs,nq,wh)
-                
-                # attention: cross-attention first
-                output, avg_attn = self.transformer_cross_attention_layers[i](
-                    output, src[level_index],
-                    memory_mask=attn_mask,
-                    memory_key_padding_mask=None,  # here we do not apply masking on padded region
-                    pos=pos[level_index], query_pos=query_embed
-                )
-
-                output = self.transformer_self_attention_layers[i](
-                    output, tgt_mask=self_tgt_mask,
-                    tgt_key_padding_mask=None,
-                    query_pos=query_embed
-                )
-                
-                # FFN
-                output = self.transformer_ffn_layers[i](
-                    output
-                )
-
-                results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels], layer_id=i, task=task)
-                attn_mask = results["attn_mask"]
-            
-            pred_captions_gen = results['outputs_captionting']
-            # pred_captions_gen = (pred_captions_gen / pred_captions_gen.norm(dim=-1, keepdim=True) + 1e-7)
-            pred_captions_gen = pred_captions_gen @ token_embs.t()
-            caping_lang_token[:,cap_idx+1] = pred_captions_gen[:,cap_idx].max(-1)[1]
-            
-        texts = self.lang_encoder.tokenizer.batch_decode(caping_lang_token, skip_special_tokens=False)
-        texts_new = []
-        
-        for x in texts:
-            x = x.split('<|endoftext|>')[0]
-            x = x.replace('<|endoftext|>','')
-            x = x.replace('<|startoftext|>','')
-            x = x.strip()
-            texts_new.append(x)
-
-        out = {'pred_captionings': caping_lang_token,
-               'pred_texts': texts_new}
-        return out
-
-
-    def forward_prediction_heads(self, output, mask_features, attn_mask_target_size, layer_id=-1, task='seg'):
-        decoder_output = self.decoder_norm(output)
-        decoder_output = decoder_output.transpose(0, 1)
-
-        # extract image captioning token from decoder output.
-        if self.task_switch['captioning'] and (task == 'vlp' or task == 'captioning_infer'):
-            outputs_captionting = decoder_output[:,self.num_queries:] @ self.caping_embed
-        else:
-            outputs_captionting = None
-
-        # recompute class token output.
-        norm_decoder_output = decoder_output / (decoder_output.norm(dim=-1, keepdim=True) + 1e-7)
-        obj_token = norm_decoder_output[:,:self.num_queries-1]
-        cls_token = norm_decoder_output[:,self.num_queries-1:self.num_queries]
-
-        sim = (cls_token @ obj_token.transpose(1,2)).softmax(-1)[:,0,:,None] # TODO include class token.
-        cls_token = (sim * decoder_output[:,:self.num_queries-1]).sum(dim=1, keepdim=True)
-
-        if (((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding']) \
-                or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-            decoder_output = torch.cat((decoder_output[:,:self.num_queries-1], cls_token, decoder_output[:,self.num_queries:2*self.num_queries-1]), dim=1)
-        else:
-            decoder_output = torch.cat((decoder_output[:,:self.num_queries-1], cls_token), dim=1)
-
-        # compute class, mask and bbox.
-        class_embed = decoder_output @ self.class_embed
-        # HACK do not compute similarity if mask is not on
-        outputs_class = self.lang_encoder.compute_similarity(class_embed, fake=(((not self.task_switch['mask']) and self.training) or (task == 'openimage')))
-
-        if self.task_switch['mask'] or self.task_switch['openimage']['mask']:
-            mask_embed = self.mask_embed(decoder_output)
-            outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)
-
-            # NOTE: prediction is of higher-resolution
-            # [B, Q, H, W] -> [B, Q, H*W] -> [B, h, Q, H*W] -> [B*h, Q, HW]
-            attn_mask = F.interpolate(outputs_mask, size=attn_mask_target_size, mode="bilinear", align_corners=False)
-
-            # must use bool type
-            # If a BoolTensor is provided, positions with ``True`` are not allowed to attend while ``False`` values will be unchanged.
-            attn_mask = (attn_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5).bool()
-            attn_mask = attn_mask.detach()
-
-            # NOTE: fill False for cls token (JY)
-            attn_mask[:, self.num_queries:self.num_queries+1].fill_(False)
-        else:
-            outputs_mask = None
-            attn_mask = torch.zeros((list(decoder_output.shape[:2]) + [attn_mask_target_size[0]*attn_mask_target_size[1]]), device=decoder_output.device).repeat(self.num_heads, 1, 1).bool()
-
-        outputs_bbox = [None for i in range(len(decoder_output))]
-        if self.task_switch['bbox']:
-            outputs_bbox = self.bbox_embed(decoder_output)
-
-        outputs_caption = None
-        if self.task_switch['caption']:
-            outputs_caption = class_embed
-            
-
-        results = {
-            "outputs_class": outputs_class,
-            "outputs_mask": outputs_mask,
-            "outputs_bbox": outputs_bbox,
-            "attn_mask": attn_mask,
-            "outputs_caption": outputs_caption,
-            "outputs_captionting": outputs_captionting,
-        }
-        return results
-
-    @torch.jit.unused
-    def _set_aux_loss(self, outputs_class, outputs_seg_masks, outputs_boxes, outputs_captions):
-        # this is a workaround to make torchscript happy, as torchscript
-        # doesn't support dictionary with non-homogeneous values, such
-        # as a dict having both a Tensor and a list.
-        if self.mask_classification:
-            return [
-                {"pred_logits": a, "pred_masks": b, "pred_boxes": c, "pred_captions": d}
-                for a, b, c, d in zip(outputs_class[:-1], outputs_seg_masks[:-1], outputs_boxes[:-1], outputs_captions[:-1])
-            ]
-        else:
-            return [{"pred_masks": b} for b in outputs_seg_masks[:-1]]
-
-
-@register_decoder
-def get_masked_transformer_decoder(cfg, in_channels, lang_encoder, mask_classification, extra):
-    return MultiScaleMaskedTransformerDecoder(cfg, in_channels, lang_encoder, mask_classification, extra)

xdecoder/body/decoder/xdecoder2.py

@@ -1,700 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-# Modified by Bowen Cheng from: https://github.com/facebookresearch/detr/blob/master/models/detr.py
-
-# --------------------------------------------------------
-# X-Decoder -- Generalized Decoding for Pixel, Image, and Language
-# Copyright (c) 2022 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Xueyan Zou (xueyan@cs.wisc.edu), Jianwei Yang (jianwyan@microsoft.com)
-# --------------------------------------------------------
-
-
-import logging
-from typing import Optional
-
-import torch
-from torch import nn, Tensor
-from torch.nn import functional as F
-
-from timm.models.layers import trunc_normal_
-from detectron2.layers import Conv2d
-import fvcore.nn.weight_init as weight_init
-
-from .registry import register_decoder
-from ...utils import configurable
-from ...modules import PositionEmbeddingSine
-
-
-class SelfAttentionLayer(nn.Module):
-
-    def __init__(self, d_model, nhead, dropout=0.0,
-                 activation="relu", normalize_before=False):
-        super().__init__()
-        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
-
-        self.norm = nn.LayerNorm(d_model)
-        self.dropout = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-        self._reset_parameters()
-    
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(self, tgt,
-                     tgt_mask: Optional[Tensor] = None,
-                     tgt_key_padding_mask: Optional[Tensor] = None,
-                     query_pos: Optional[Tensor] = None):
-        q = k = self.with_pos_embed(tgt, query_pos)
-        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
-                              key_padding_mask=tgt_key_padding_mask)[0]
-        tgt = tgt + self.dropout(tgt2)
-        tgt = self.norm(tgt)
-
-        return tgt
-
-    def forward_pre(self, tgt,
-                    tgt_mask: Optional[Tensor] = None,
-                    tgt_key_padding_mask: Optional[Tensor] = None,
-                    query_pos: Optional[Tensor] = None):
-        tgt2 = self.norm(tgt)
-        q = k = self.with_pos_embed(tgt2, query_pos)
-        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
-                              key_padding_mask=tgt_key_padding_mask)[0]
-        tgt = tgt + self.dropout(tgt2)
-        
-        return tgt
-
-    def forward(self, tgt,
-                tgt_mask: Optional[Tensor] = None,
-                tgt_key_padding_mask: Optional[Tensor] = None,
-                query_pos: Optional[Tensor] = None):
-        if self.normalize_before:
-            return self.forward_pre(tgt, tgt_mask,
-                                    tgt_key_padding_mask, query_pos)
-        return self.forward_post(tgt, tgt_mask,
-                                 tgt_key_padding_mask, query_pos)
-
-
-class CrossAttentionLayer(nn.Module):
-
-    def __init__(self, d_model, nhead, dropout=0.0,
-                 activation="relu", normalize_before=False):
-        super().__init__()
-        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
-
-        self.norm = nn.LayerNorm(d_model)
-        self.dropout = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-        self._reset_parameters()
-    
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(self, tgt, memory,
-                     memory_mask: Optional[Tensor] = None,
-                     memory_key_padding_mask: Optional[Tensor] = None,
-                     pos: Optional[Tensor] = None,
-                     query_pos: Optional[Tensor] = None):
-        tgt2, avg_attn = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
-                                   key=self.with_pos_embed(memory, pos),
-                                   value=memory, attn_mask=memory_mask,
-                                   key_padding_mask=memory_key_padding_mask)
-        tgt = tgt + self.dropout(tgt2)
-        tgt = self.norm(tgt)
-        return tgt, avg_attn
-
-    def forward_pre(self, tgt, memory,
-                    memory_mask: Optional[Tensor] = None,
-                    memory_key_padding_mask: Optional[Tensor] = None,
-                    pos: Optional[Tensor] = None,
-                    query_pos: Optional[Tensor] = None):
-        tgt2 = self.norm(tgt)
-        tgt2, avg_attn = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
-                                   key=self.with_pos_embed(memory, pos),
-                                   value=memory, attn_mask=memory_mask,
-                                   key_padding_mask=memory_key_padding_mask)
-        tgt = tgt + self.dropout(tgt2)
-
-        return tgt, avg_attn
-
-    def forward(self, tgt, memory,
-                memory_mask: Optional[Tensor] = None,
-                memory_key_padding_mask: Optional[Tensor] = None,
-                pos: Optional[Tensor] = None,
-                query_pos: Optional[Tensor] = None):
-        if self.normalize_before:
-            return self.forward_pre(tgt, memory, memory_mask,
-                                    memory_key_padding_mask, pos, query_pos)
-        return self.forward_post(tgt, memory, memory_mask,
-                                 memory_key_padding_mask, pos, query_pos)
-
-
-class FFNLayer(nn.Module):
-
-    def __init__(self, d_model, dim_feedforward=2048, dropout=0.0,
-                 activation="relu", normalize_before=False):
-        super().__init__()
-        # Implementation of Feedforward model
-        self.linear1 = nn.Linear(d_model, dim_feedforward)
-        self.dropout = nn.Dropout(dropout)
-        self.linear2 = nn.Linear(dim_feedforward, d_model)
-
-        self.norm = nn.LayerNorm(d_model)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-        self._reset_parameters()
-    
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(self, tgt):
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
-        tgt = tgt + self.dropout(tgt2)
-        tgt = self.norm(tgt)
-        return tgt
-
-    def forward_pre(self, tgt):
-        tgt2 = self.norm(tgt)
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
-        tgt = tgt + self.dropout(tgt2)
-        return tgt
-
-    def forward(self, tgt):
-        if self.normalize_before:
-            return self.forward_pre(tgt)
-        return self.forward_post(tgt)
-
-
-def _get_activation_fn(activation):
-    """Return an activation function given a string"""
-    if activation == "relu":
-        return F.relu
-    if activation == "gelu":
-        return F.gelu
-    if activation == "glu":
-        return F.glu
-    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
-
-
-class MLP(nn.Module):
-    """ Very simple multi-layer perceptron (also called FFN)"""
-
-    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
-        super().__init__()
-        self.num_layers = num_layers
-        h = [hidden_dim] * (num_layers - 1)
-        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
-
-    def forward(self, x):
-        for i, layer in enumerate(self.layers):
-            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
-        return x
-
-
-class MultiScaleMaskedTransformerDecoder(nn.Module):
-
-    _version = 2
-
-    @configurable
-    def __init__(
-        self,
-        lang_encoder: nn.Module,
-        in_channels,
-        mask_classification=True,
-        *,
-        hidden_dim: int,
-        dim_proj: int,
-        num_queries: int,
-        contxt_len: int,
-        nheads: int,
-        dim_feedforward: int,
-        dec_layers: int,
-        pre_norm: bool,
-        mask_dim: int,
-        task_switch: dict,
-        captioning_step: int,
-        enforce_input_project: bool,
-    ):
-        """
-        NOTE: this interface is experimental.
-        Args:
-            in_channels: channels of the input features
-            mask_classification: whether to add mask classifier or not
-            num_classes: number of classes
-            hidden_dim: Transformer feature dimension
-            num_queries: number of queries
-            nheads: number of heads
-            dim_feedforward: feature dimension in feedforward network
-            enc_layers: number of Transformer encoder layers
-            dec_layers: number of Transformer decoder layers
-            pre_norm: whether to use pre-LayerNorm or not
-            mask_dim: mask feature dimension
-            enforce_input_project: add input project 1x1 conv even if input
-                channels and hidden dim is identical
-        """
-        super().__init__()
-        assert mask_classification, "Only support mask classification model"
-        self.mask_classification = mask_classification
-
-        # positional encoding
-        N_steps = hidden_dim // 2
-        self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
-        
-        # define Transformer decoder here
-        self.num_heads = nheads
-        self.num_layers = dec_layers
-        self.contxt_len = contxt_len
-        self.transformer_self_attention_layers = nn.ModuleList()
-        self.transformer_cross_attention_layers = nn.ModuleList()
-        self.transformer_ffn_layers = nn.ModuleList()
-
-        for _ in range(self.num_layers):
-            self.transformer_self_attention_layers.append(
-                SelfAttentionLayer(
-                    d_model=hidden_dim,
-                    nhead=nheads,
-                    dropout=0.0,
-                    normalize_before=pre_norm,
-                )
-            )
-
-            self.transformer_cross_attention_layers.append(
-                CrossAttentionLayer(
-                    d_model=hidden_dim,
-                    nhead=nheads,
-                    dropout=0.0,
-                    normalize_before=pre_norm,
-                )
-            )
-
-            self.transformer_ffn_layers.append(
-                FFNLayer(
-                    d_model=hidden_dim,
-                    dim_feedforward=dim_feedforward,
-                    dropout=0.0,
-                    normalize_before=pre_norm,
-                )
-            )
-
-        self.decoder_norm = nn.LayerNorm(hidden_dim)
-
-        self.num_queries = num_queries
-        # learnable query features
-        self.query_feat = nn.Embedding(num_queries, hidden_dim)
-        # learnable query p.e.
-        self.query_embed = nn.Embedding(num_queries, hidden_dim)
-        
-        # level embedding (we always use 3 scales)
-        self.num_feature_levels = 3
-        self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim)
-        self.input_proj = nn.ModuleList()
-        
-        for _ in range(self.num_feature_levels):
-            if in_channels != hidden_dim or enforce_input_project:
-                self.input_proj.append(Conv2d(in_channels, hidden_dim, kernel_size=1))
-                weight_init.c2_xavier_fill(self.input_proj[-1])
-            else:
-                self.input_proj.append(nn.Sequential())
-
-        self.task_switch = task_switch
-
-        # output FFNs
-        self.lang_encoder = lang_encoder
-        if self.task_switch['mask']:
-            self.mask_embed = MLP(hidden_dim, hidden_dim, mask_dim, 3)
-
-        self.class_embed = nn.Parameter(torch.empty(hidden_dim, dim_proj))
-        trunc_normal_(self.class_embed, std=.02)
-
-        if task_switch['bbox']:
-            self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
-
-        # Caption Project and query
-        if task_switch['captioning']:
-            self.caping_embed = nn.Parameter(torch.empty(hidden_dim, dim_proj))
-            trunc_normal_(self.caping_embed, std=.02)
-            self.query_feat_caping = nn.Embedding(contxt_len, hidden_dim)
-            # self.pos_embed_caping = nn.Embedding(contxt_len, hidden_dim)
-            self.captioning_step = captioning_step
-
-        # register self_attn_mask to avoid information leakage, it includes interaction between object query, class query and caping query
-        self_attn_mask = torch.zeros((1, num_queries + contxt_len, num_queries + contxt_len)).bool()
-        self_attn_mask[:, :num_queries, num_queries:] = True # object+class query does not attend with caption query.
-        self_attn_mask[:, num_queries:, num_queries:] = torch.triu(torch.ones((1, contxt_len, contxt_len)), diagonal=1).bool() # caption query only attend with previous token.
-        self_attn_mask[:, :num_queries-1, num_queries-1:num_queries] = True # object query does not attend with class query.
-        self_attn_mask[:, num_queries-1:num_queries, :num_queries-1] = True # class query does not attend with object query.
-        self.register_buffer("self_attn_mask", self_attn_mask)
-
-
-    @classmethod
-    def from_config(cls, cfg, in_channels, lang_encoder, mask_classification, extra):
-        ret = {}
-
-        ret["lang_encoder"] = lang_encoder
-        ret["in_channels"] = in_channels
-        ret["mask_classification"] = mask_classification
-
-        enc_cfg = cfg['MODEL']['ENCODER']
-        dec_cfg = cfg['MODEL']['DECODER']
-        
-        ret["hidden_dim"] = dec_cfg['HIDDEN_DIM']
-        ret["dim_proj"] = cfg['MODEL']['DIM_PROJ']
-        ret["num_queries"] = dec_cfg['NUM_OBJECT_QUERIES']
-        ret["contxt_len"] = cfg['MODEL']['TEXT']['CONTEXT_LENGTH']
-        
-        # Transformer parameters:
-        ret["nheads"] = dec_cfg['NHEADS']
-        ret["dim_feedforward"] = dec_cfg['DIM_FEEDFORWARD']
-
-        # NOTE: because we add learnable query features which requires supervision,
-        # we add minus 1 to decoder layers to be consistent with our loss
-        # implementation: that is, number of auxiliary losses is always
-        # equal to number of decoder layers. With learnable query features, the number of
-        # auxiliary losses equals number of decoders plus 1.
-        assert dec_cfg['DEC_LAYERS'] >= 1
-        ret["dec_layers"] = dec_cfg['DEC_LAYERS'] - 1
-        ret["pre_norm"] = dec_cfg['PRE_NORM']
-        ret["enforce_input_project"] = dec_cfg['ENFORCE_INPUT_PROJ']
-        ret["mask_dim"] = enc_cfg['MASK_DIM']
-
-        ret["task_switch"] = extra['task_switch']
-        ret["captioning_step"] = dec_cfg['CAPTIONING'].get('STEP', 50)
-
-        return ret
-
-    def forward(self, x, mask_features, mask=None, target_queries=None, target_vlp=None, task='seg', extra={}):
-        if task == 'captioning_infer':
-            return self.forward_captioning(x, mask_features, mask=mask, target_queries=target_queries, target_vlp=target_vlp, task=task, extra=extra)
-        # x is a list of multi-scale feature
-        assert len(x) == self.num_feature_levels
-        src = []
-        pos = []
-        size_list = []
-        
-        # disable mask, it does not affect performance
-        del mask
-        for i in range(self.num_feature_levels):
-            size_list.append(x[i].shape[-2:])
-            pos.append(self.pe_layer(x[i], None).flatten(2))
-            src.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
-
-            # flatten NxCxHxW to HWxNxC
-            pos[-1] = pos[-1].permute(2, 0, 1)
-            src[-1] = src[-1].permute(2, 0, 1)
-
-        _, bs, _ = src[0].shape
-
-        # QxNxC
-        query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
-        output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)
-
-        predictions_class = []
-        predictions_mask = []
-        predictions_bbox = []
-        predictions_caption = []
-        predictions_captioning = []
-        
-        self_tgt_mask = None
-        if self.training and task == 'vlp' and self.task_switch['captioning']:
-            output = torch.cat((output, self.query_feat_caping.weight.unsqueeze(1).repeat(1, bs, 1)), dim=0) # concat object query, class token and caption token.
-            caping_lang_embed = torch.cat([caption['caption_tokens'] for caption in target_vlp], dim=0).transpose(0, 1) # language output
-            # _caping_lang_embed = caping_lang_embed.detach().clone()
-            # output = torch.cat((output, _caping_lang_embed), dim=0) # concat object query, class token and caption token.
-            # caping_lang_embed += self.pos_embed_caping.weight.unsqueeze(1).repeat(1, bs, 1)
-            query_embed = torch.cat((query_embed, caping_lang_embed), dim=0) # may not add at the beginning.
-            self_tgt_mask = self.self_attn_mask.repeat(output.shape[1]*self.num_heads, 1, 1)
-        elif (((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding']) \
-                or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-            self_tgt_mask = self.self_attn_mask[:,:self.num_queries,:self.num_queries].repeat(output.shape[1]*self.num_heads, 1, 1)
-            grounding_tokens = extra['grounding_tokens']
-            _grounding_tokens = grounding_tokens.detach().clone()
-            # initialize with negative attention at the beginning.
-            pad_tgt_mask = torch.ones((1, self.num_queries + (self.num_queries-1) + len(grounding_tokens), self.num_queries + (self.num_queries-1) + len(grounding_tokens)), device=self_tgt_mask.device).bool().repeat(output.shape[1]*self.num_heads, 1, 1)
-            pad_tgt_mask[:,:self.num_queries,:self.num_queries] = self_tgt_mask
-            pad_tgt_mask[:,self.num_queries:,self.num_queries:] = False # grounding tokens could attend with eatch other
-            self_tgt_mask = pad_tgt_mask
-            output = torch.cat((output, output[:-1]), dim=0)
-            query_embed = torch.cat((query_embed, query_embed[:-1]), dim=0) # also pad language embdding to fix embedding
-        else:
-            self_tgt_mask = self.self_attn_mask[:,:self.num_queries,:self.num_queries].repeat(output.shape[1]*self.num_heads, 1, 1)
-
-        # prediction heads on learnable query features
-        results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[0], task=task)
-        attn_mask = results["attn_mask"]
-        predictions_class.append(results["outputs_class"])
-        predictions_mask.append(results["outputs_mask"])
-        predictions_bbox.append(results["outputs_bbox"])
-        predictions_caption.append(results["outputs_caption"])
-        predictions_captioning.append(results["outputs_captionting"])
-        
-        for i in range(self.num_layers):
-            level_index = i % self.num_feature_levels
-            attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
-
-            if self.training and task == 'vlp' and self.task_switch['captioning']:
-                attn_mask = torch.cat((attn_mask, torch.zeros_like(attn_mask[:, :self.contxt_len, :])), dim=1)
-            # attention: cross-attention first
-            output, avg_attn = self.transformer_cross_attention_layers[i](
-                output, src[level_index],
-                memory_mask=attn_mask,
-                memory_key_padding_mask=None,  # here we do not apply masking on padded region
-                pos=pos[level_index], query_pos=query_embed
-            )
-
-            if (((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding']) \
-                    or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-                output = torch.cat((output, _grounding_tokens), dim=0)
-                query_embed = torch.cat((query_embed, grounding_tokens), dim=0)
-
-            output = self.transformer_self_attention_layers[i](
-                output, tgt_mask=self_tgt_mask,
-                tgt_key_padding_mask=None,
-                query_pos=query_embed
-            )
-            
-            # FFN
-            output = self.transformer_ffn_layers[i](
-                output
-            )
-
-            if ((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding'] \
-                    or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-                _grounding_tokens = output[-len(_grounding_tokens):]
-                output = output[:-len(_grounding_tokens)]
-                query_embed = query_embed[:-len(_grounding_tokens)]
-
-            results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels], layer_id=i, task=task)
-            attn_mask = results["attn_mask"]
-            predictions_class.append(results["outputs_class"])
-            predictions_mask.append(results["outputs_mask"])
-            predictions_bbox.append(results["outputs_bbox"])
-            predictions_caption.append(results["outputs_caption"])
-            predictions_captioning.append(results["outputs_captionting"])
-
-        assert len(predictions_class) == self.num_layers + 1
-        if task == 'vlp':
-            out = {'pred_captionings': predictions_captioning[-1], 
-                   'pred_captions': predictions_caption[-1], 
-                   'aux_outputs': [{'pred_captionings': x, 'pred_captions': y } for x, y in zip(predictions_captioning[:-1], predictions_caption[:-1])]}
-            return out
-        else:
-            out = {
-                'pred_logits': predictions_class[-1],
-                'pred_masks': predictions_mask[-1],
-                'pred_boxes': predictions_bbox[-1],
-                'pred_captions': predictions_caption[-1],
-                'aux_outputs': self._set_aux_loss(
-                    predictions_class if self.mask_classification else None, predictions_mask, predictions_bbox, predictions_caption
-                )
-            }
-            return out
-
-    def forward_captioning(self, x, mask_features, mask = None, target_queries = None, target_vlp = None, task='seg', extra={}):
-        # x is a list of multi-scale feature
-        assert len(x) == self.num_feature_levels
-        src = []
-        pos = []
-        size_list = []
-        
-        # disable mask, it does not affect performance
-        del mask
-        for i in range(self.num_feature_levels):
-            size_list.append(x[i].shape[-2:])
-            pos.append(self.pe_layer(x[i], None).flatten(2))
-            src.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
-
-            # flatten NxCxHxW to HWxNxC
-            pos[-1] = pos[-1].permute(2, 0, 1)
-            src[-1] = src[-1].permute(2, 0, 1)
-
-        _, bs, _ = src[0].shape
-
-        # QxNxC
-        query_embed_ = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
-        query_feat = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)        
-        caping_lang_token = extra['start_token'].repeat(bs, 1)
-        start_id = 0
-        if 'token' in extra:
-            caping_lang_token[:,:len(extra['token'][0])] = extra['token']
-            start_id = len(extra['token'][0])-1
-        query_feat_caping = self.query_feat_caping.weight.unsqueeze(1).repeat(1, bs, 1)
-        # pos_embed_caping = self.pos_embed_caping.weight.unsqueeze(1).repeat(1, bs, 1)
-        # prepare token embedding for evaluation
-        token_embs = self.lang_encoder.lang_encoder.token_embedding.weight
-        # token_embs = (token_embs / token_embs.norm(dim=-1, keepdim=True) + 1e-7)
-        
-        for cap_idx in range(start_id, self.captioning_step):
-            caping_lang_embed = self.lang_encoder.forward_language_token((caping_lang_token,))[0].transpose(0, 1)
-            # output = torch.cat((query_feat, caping_lang_embed), dim=0) # concat object query, class token and caption token.
-            # caping_lang_embed += pos_embed_caping
-            query_embed = torch.cat((query_embed_, caping_lang_embed), dim=0) # may not add at the beginning.
-            output = torch.cat((query_feat, query_feat_caping), dim=0) # concat object query, class token and caption token.
-
-            # prediction heads on learnable query features
-            results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[0], task=task)
-            attn_mask = results["attn_mask"]
-        
-            for i in range(self.num_layers):
-                level_index = i % self.num_feature_levels
-                attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
-                attn_mask = torch.cat((attn_mask, torch.zeros_like(attn_mask[:, :self.contxt_len, :])), dim=1)
-                self_tgt_mask = self.self_attn_mask.repeat(output.shape[1]*self.num_heads, 1, 1)
-
-                if extra['captioning_mask'] is not None:
-                    bs,nq,wh = attn_mask.shape
-                    assert bs==self.num_heads, "Only support single image referring captioning."
-                    cap_mask = extra['captioning_mask']
-                    attn_mask = attn_mask.reshape(bs,nq,size_list[i%3][0],size_list[i%3][1])
-                    cap_mask = F.interpolate(cap_mask[None,].float(), size_list[i%3], mode='nearest').bool()[0,0]
-                    attn_mask[:,self.num_queries:, cap_mask] = True
-                    attn_mask = attn_mask.reshape(bs,nq,wh)
-                
-                # attention: cross-attention first
-                output, avg_attn = self.transformer_cross_attention_layers[i](
-                    output, src[level_index],
-                    memory_mask=attn_mask,
-                    memory_key_padding_mask=None,  # here we do not apply masking on padded region
-                    pos=pos[level_index], query_pos=query_embed
-                )
-
-                output = self.transformer_self_attention_layers[i](
-                    output, tgt_mask=self_tgt_mask,
-                    tgt_key_padding_mask=None,
-                    query_pos=query_embed
-                )
-                
-                # FFN
-                output = self.transformer_ffn_layers[i](
-                    output
-                )
-
-                results = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels], layer_id=i, task=task)
-                attn_mask = results["attn_mask"]
-            
-            pred_captions_gen = results['outputs_captionting']
-            # pred_captions_gen = (pred_captions_gen / pred_captions_gen.norm(dim=-1, keepdim=True) + 1e-7)
-            pred_captions_gen = pred_captions_gen @ token_embs.t()
-            caping_lang_token[:,cap_idx+1] = pred_captions_gen[:,cap_idx].max(-1)[1]
-            
-        texts = self.lang_encoder.tokenizer.batch_decode(caping_lang_token, skip_special_tokens=False)
-        texts_new = []
-        
-        for x in texts:
-            x = x.split('<|endoftext|>')[0]
-            x = x.replace('<|endoftext|>','')
-            x = x.replace('<|startoftext|>','')
-            x = x.strip()
-            texts_new.append(x)
-
-        out = {'pred_captionings': caping_lang_token,
-               'pred_texts': texts_new}
-        return out
-
-
-    def forward_prediction_heads(self, output, mask_features, attn_mask_target_size, layer_id=-1, task='seg'):
-        decoder_output = self.decoder_norm(output)
-        decoder_output = decoder_output.transpose(0, 1)
-
-        # extract image captioning token from decoder output.
-        if self.task_switch['captioning'] and (task == 'vlp' or task == 'captioning_infer'):
-            outputs_captionting = decoder_output[:,self.num_queries:] @ self.caping_embed
-        else:
-            outputs_captionting = None
-
-        # recompute class token output.
-        norm_decoder_output = decoder_output / (decoder_output.norm(dim=-1, keepdim=True) + 1e-7)
-        obj_token = norm_decoder_output[:,:self.num_queries-1]
-        cls_token = norm_decoder_output[:,self.num_queries-1:self.num_queries]
-
-        sim = (cls_token @ obj_token.transpose(1,2)).softmax(-1)[:,0,:,None] # TODO include class token.
-        cls_token = (sim * decoder_output[:,:self.num_queries-1]).sum(dim=1, keepdim=True)
-
-        if (((self.training and task == 'seg') or (task == 'grounding_eval')) and self.task_switch['grounding']) \
-                or ((self.training and task == 'openimage') and self.task_switch['openimage']['grounding']):
-            decoder_output = torch.cat((decoder_output[:,:self.num_queries-1], cls_token, decoder_output[:,self.num_queries:2*self.num_queries-1]), dim=1)
-        else:
-            decoder_output = torch.cat((decoder_output[:,:self.num_queries-1], cls_token), dim=1)
-
-        # compute class, mask and bbox.
-        class_embed = decoder_output @ self.class_embed
-        # HACK do not compute similarity if mask is not on
-        outputs_class = self.lang_encoder.compute_similarity(class_embed, fake=(((not self.task_switch['mask']) and self.training) or (task == 'openimage')))
-
-        if self.task_switch['mask'] or self.task_switch['openimage']['mask']:
-            mask_embed = self.mask_embed(decoder_output)
-            outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)
-
-            # NOTE: prediction is of higher-resolution
-            # [B, Q, H, W] -> [B, Q, H*W] -> [B, h, Q, H*W] -> [B*h, Q, HW]
-            attn_mask = F.interpolate(outputs_mask, size=attn_mask_target_size, mode="bilinear", align_corners=False)
-
-            # must use bool type
-            # If a BoolTensor is provided, positions with ``True`` are not allowed to attend while ``False`` values will be unchanged.
-            attn_mask = (attn_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5).bool()
-            attn_mask = attn_mask.detach()
-
-            # NOTE: fill False for cls token (JY)
-            attn_mask[:, self.num_queries:self.num_queries+1].fill_(False)
-        else:
-            outputs_mask = None
-            attn_mask = torch.zeros((list(decoder_output.shape[:2]) + [attn_mask_target_size[0]*attn_mask_target_size[1]]), device=decoder_output.device).repeat(self.num_heads, 1, 1).bool()
-
-        outputs_bbox = [None for i in range(len(decoder_output))]
-        if self.task_switch['bbox']:
-            outputs_bbox = self.bbox_embed(decoder_output)
-
-        outputs_caption = None
-        if self.task_switch['caption']:
-            outputs_caption = class_embed
-            
-
-        results = {
-            "outputs_class": outputs_class,
-            "outputs_mask": outputs_mask,
-            "outputs_bbox": outputs_bbox,
-            "attn_mask": attn_mask,
-            "outputs_caption": outputs_caption,
-            "outputs_captionting": outputs_captionting,
-        }
-        return results
-
-    @torch.jit.unused
-    def _set_aux_loss(self, outputs_class, outputs_seg_masks, outputs_boxes, outputs_captions):
-        # this is a workaround to make torchscript happy, as torchscript
-        # doesn't support dictionary with non-homogeneous values, such
-        # as a dict having both a Tensor and a list.
-        if self.mask_classification:
-            return [
-                {"pred_logits": a, "pred_masks": b, "pred_boxes": c, "pred_captions": d}
-                for a, b, c, d in zip(outputs_class[:-1], outputs_seg_masks[:-1], outputs_boxes[:-1], outputs_captions[:-1])
-            ]
-        else:
-            return [{"pred_masks": b} for b in outputs_seg_masks[:-1]]
-
-
-@register_decoder
-def get_masked_transformer_decoder(cfg, in_channels, lang_encoder, mask_classification, extra):
-    return MultiScaleMaskedTransformerDecoder(cfg, in_channels, lang_encoder, mask_classification, extra)

xdecoder/body/encoder/__init__.py

@@ -1 +0,0 @@
-from .build import build_encoder

xdecoder/body/encoder/build.py

@@ -1,12 +0,0 @@
-from .registry import model_entrypoints
-from .registry import is_model
-
-from .transformer_encoder_fpn import *
-
-def build_encoder(config, *args, **kwargs):
-    model_name = config['MODEL']['ENCODER']['NAME']
-
-    if not is_model(model_name):
-        raise ValueError(f'Unkown model: {model_name}')
-
-    return model_entrypoints(model_name)(config, *args, **kwargs)

xdecoder/body/encoder/registry.py

@@ -1,13 +0,0 @@
-_model_entrypoints = {}
-
-def register_encoder(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/body/encoder/transformer_encoder_fpn.py

@@ -1,324 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-import logging
-import numpy as np
-from typing import Callable, Dict, List, Optional, Tuple, Union
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-from torch.nn.init import xavier_uniform_, constant_, uniform_, normal_
-from torch.cuda.amp import autocast
-
-import fvcore.nn.weight_init as weight_init
-from detectron2.layers import Conv2d, DeformConv, ShapeSpec, get_norm
-
-from .registry import register_encoder
-from ..transformer_blocks import TransformerEncoder, TransformerEncoderLayer, _get_clones, _get_activation_fn
-from ...modules import PositionEmbeddingSine
-from ...utils import configurable
-
-# from ..layers import Conv2d, DeformConv, ShapeSpec, get_norm
-
-# This is a modified FPN decoder.
-class BasePixelDecoder(nn.Module):
-    def __init__(
-        self,
-        input_shape: Dict[str, ShapeSpec],
-        *,
-        conv_dim: int,
-        mask_dim: int,
-        mask_on: bool,
-        norm: Optional[Union[str, Callable]] = None,
-    ):
-        """
-        NOTE: this interface is experimental.
-        Args:
-            input_shape: shapes (channels and stride) of the input features
-            conv_dims: number of output channels for the intermediate conv layers.
-            mask_dim: number of output channels for the final conv layer.
-            norm (str or callable): normalization for all conv layers
-        """
-        super().__init__()
-
-        input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
-        self.in_features = [k for k, v in input_shape]  # starting from "res2" to "res5"
-        feature_channels = [v.channels for k, v in input_shape]
-
-        lateral_convs = []
-        output_convs = []
-
-        use_bias = norm == ""
-        for idx, in_channels in enumerate(feature_channels):
-            if idx == len(self.in_features) - 1:
-                output_norm = get_norm(norm, conv_dim)
-                output_conv = Conv2d(
-                    in_channels,
-                    conv_dim,
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                    bias=use_bias,
-                    norm=output_norm,
-                    activation=F.relu,
-                )
-                weight_init.c2_xavier_fill(output_conv)
-                self.add_module("layer_{}".format(idx + 1), output_conv)
-
-                lateral_convs.append(None)
-                output_convs.append(output_conv)
-            else:
-                lateral_norm = get_norm(norm, conv_dim)
-                output_norm = get_norm(norm, conv_dim)
-
-                lateral_conv = Conv2d(
-                    in_channels, conv_dim, kernel_size=1, bias=use_bias, norm=lateral_norm
-                )
-                output_conv = Conv2d(
-                    conv_dim,
-                    conv_dim,
-                    kernel_size=3,
-                    stride=1,
-                    padding=1,
-                    bias=use_bias,
-                    norm=output_norm,
-                    activation=F.relu,
-                )
-                weight_init.c2_xavier_fill(lateral_conv)
-                weight_init.c2_xavier_fill(output_conv)
-                self.add_module("adapter_{}".format(idx + 1), lateral_conv)
-                self.add_module("layer_{}".format(idx + 1), output_conv)
-
-                lateral_convs.append(lateral_conv)
-                output_convs.append(output_conv)
-        # Place convs into top-down order (from low to high resolution)
-        # to make the top-down computation in forward clearer.
-        self.lateral_convs = lateral_convs[::-1]
-        self.output_convs = output_convs[::-1]
-
-        self.mask_on = mask_on
-        if self.mask_on:
-            self.mask_dim = mask_dim
-            self.mask_features = Conv2d(
-                conv_dim,
-                mask_dim,
-                kernel_size=3,
-                stride=1,
-                padding=1,
-            )
-            weight_init.c2_xavier_fill(self.mask_features)
-
-        self.maskformer_num_feature_levels = 3  # always use 3 scales
-
-    @classmethod
-    def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
-        enc_cfg = cfg['MODEL']['ENCODER']
-        ret = {}
-        ret["input_shape"] = {
-            k: v for k, v in input_shape.items() if k in enc_cfg['IN_FEATURES']
-        }
-        ret["conv_dim"] = enc_cfg['CONVS_DIM']
-        ret["mask_dim"] = enc_cfg['MASK_DIM']
-        ret["norm"] = enc_cfg['NORM']
-        return ret
-
-    def forward_features(self, features):
-        multi_scale_features = []
-        num_cur_levels = 0
-        # Reverse feature maps into top-down order (from low to high resolution)
-        for idx, f in enumerate(self.in_features[::-1]):
-            x = features[f]
-            lateral_conv = self.lateral_convs[idx]
-            output_conv = self.output_convs[idx]
-            if lateral_conv is None:
-                y = output_conv(x)
-            else:
-                cur_fpn = lateral_conv(x)
-                # Following FPN implementation, we use nearest upsampling here
-                y = cur_fpn + F.interpolate(y, size=cur_fpn.shape[-2:], mode="nearest")
-                y = output_conv(y)
-            if num_cur_levels < self.maskformer_num_feature_levels:
-                multi_scale_features.append(y)
-                num_cur_levels += 1
-        
-        mask_features = self.mask_features(y) if self.mask_on else None
-        return mask_features, None, multi_scale_features
-
-    def forward(self, features, targets=None):
-        logger = logging.getLogger(__name__)
-        logger.warning("Calling forward() may cause unpredicted behavior of PixelDecoder module.")
-        return self.forward_features(features)
-
-
-class TransformerEncoderOnly(nn.Module):
-    def __init__(
-        self,
-        d_model=512,
-        nhead=8,
-        num_encoder_layers=6,
-        dim_feedforward=2048,
-        dropout=0.1,
-        activation="relu",
-        normalize_before=False,
-    ):
-        super().__init__()
-
-        encoder_layer = TransformerEncoderLayer(
-            d_model, nhead, dim_feedforward, dropout, activation, normalize_before
-        )
-        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
-        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
-
-        self._reset_parameters()
-
-        self.d_model = d_model
-        self.nhead = nhead
-
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def forward(self, src, mask, pos_embed):
-        # flatten NxCxHxW to HWxNxC
-        bs, c, h, w = src.shape
-        src = src.flatten(2).permute(2, 0, 1)
-        pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
-        if mask is not None:
-            mask = mask.flatten(1)
-
-        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
-        return memory.permute(1, 2, 0).view(bs, c, h, w)
-
-
-# This is a modified FPN decoder with extra Transformer encoder that processes the lowest-resolution feature map.
-class TransformerEncoderPixelDecoder(BasePixelDecoder):
-    @configurable
-    def __init__(
-        self,
-        input_shape: Dict[str, ShapeSpec],
-        *,
-        transformer_dropout: float,
-        transformer_nheads: int,
-        transformer_dim_feedforward: int,
-        transformer_enc_layers: int,
-        transformer_pre_norm: bool,
-        conv_dim: int,
-        mask_dim: int,
-        mask_on: int,
-        norm: Optional[Union[str, Callable]] = None,
-    ):
-        """
-        NOTE: this interface is experimental.
-        Args:
-            input_shape: shapes (channels and stride) of the input features
-            transformer_dropout: dropout probability in transformer
-            transformer_nheads: number of heads in transformer
-            transformer_dim_feedforward: dimension of feedforward network
-            transformer_enc_layers: number of transformer encoder layers
-            transformer_pre_norm: whether to use pre-layernorm or not
-            conv_dims: number of output channels for the intermediate conv layers.
-            mask_dim: number of output channels for the final conv layer.
-            norm (str or callable): normalization for all conv layers
-        """
-        super().__init__(input_shape, conv_dim=conv_dim, mask_dim=mask_dim, norm=norm, mask_on=mask_on)
-
-        input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
-        self.in_features = [k for k, v in input_shape]  # starting from "res2" to "res5"
-        feature_strides = [v.stride for k, v in input_shape]
-        feature_channels = [v.channels for k, v in input_shape]
-
-        in_channels = feature_channels[len(self.in_features) - 1]
-        self.input_proj = Conv2d(in_channels, conv_dim, kernel_size=1)
-        weight_init.c2_xavier_fill(self.input_proj)
-        self.transformer = TransformerEncoderOnly(
-            d_model=conv_dim,
-            dropout=transformer_dropout,
-            nhead=transformer_nheads,
-            dim_feedforward=transformer_dim_feedforward,
-            num_encoder_layers=transformer_enc_layers,
-            normalize_before=transformer_pre_norm,
-        )
-        N_steps = conv_dim // 2
-        self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
-
-        # update layer
-        use_bias = norm == ""
-        output_norm = get_norm(norm, conv_dim)
-        output_conv = Conv2d(
-            conv_dim,
-            conv_dim,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            bias=use_bias,
-            norm=output_norm,
-            activation=F.relu,
-        )
-        weight_init.c2_xavier_fill(output_conv)
-        delattr(self, "layer_{}".format(len(self.in_features)))
-        self.add_module("layer_{}".format(len(self.in_features)), output_conv)
-        self.output_convs[0] = output_conv
-
-    @classmethod
-    def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
-        enc_cfg = cfg['MODEL']['ENCODER']
-        dec_cfg = cfg['MODEL']['DECODER']
-
-        ret = super().from_config(cfg, input_shape)
-        ret["transformer_dropout"] = dec_cfg['DROPOUT']
-        ret["transformer_nheads"] = dec_cfg['NHEADS']
-        ret["transformer_dim_feedforward"] = dec_cfg['DIM_FEEDFORWARD']
-        ret["transformer_enc_layers"] = enc_cfg['TRANSFORMER_ENC_LAYERS']  # a separate config
-        ret["transformer_pre_norm"] = dec_cfg['PRE_NORM']
-
-        ret['mask_on'] = cfg['MODEL']['DECODER']['MASK']
-        return ret
-
-    def forward_features(self, features):
-        multi_scale_features = []
-        num_cur_levels = 0
-        
-        # Reverse feature maps into top-down order (from low to high resolution)
-        for idx, f in enumerate(self.in_features[::-1]):
-            x = features[f]
-            lateral_conv = self.lateral_convs[idx]
-            output_conv = self.output_convs[idx]
-            if lateral_conv is None:
-                transformer = self.input_proj(x)
-                pos = self.pe_layer(x)
-                transformer = self.transformer(transformer, None, pos)
-                y = output_conv(transformer)
-                # save intermediate feature as input to Transformer decoder
-                transformer_encoder_features = transformer
-            else:
-                cur_fpn = lateral_conv(x)
-                # Following FPN implementation, we use nearest upsampling here
-                y = cur_fpn + F.interpolate(y, size=cur_fpn.shape[-2:], mode="nearest")
-                y = output_conv(y)
-            if num_cur_levels < self.maskformer_num_feature_levels:
-                multi_scale_features.append(y)
-                num_cur_levels += 1
-
-        mask_features = self.mask_features(y) if self.mask_on else None
-        return mask_features, transformer_encoder_features, multi_scale_features
-
-    def forward(self, features, targets=None):
-        logger = logging.getLogger(__name__)
-        logger.warning("Calling forward() may cause unpredicted behavior of PixelDecoder module.")
-        return self.forward_features(features)
-
-
-
-@register_encoder
-def get_transformer_encoder_fpn(cfg, input_shape):
-    """
-    Build a pixel decoder from `cfg.MODEL.MASK_FORMER.PIXEL_DECODER_NAME`.
-    """
-    model = TransformerEncoderPixelDecoder(cfg, input_shape)    
-    forward_features = getattr(model, "forward_features", None)
-    if not callable(forward_features):
-        raise ValueError(
-            "Only SEM_SEG_HEADS with forward_features method can be used as pixel decoder. "
-            f"Please implement forward_features for {name} to only return mask features."
-        )
-    return model

xdecoder/body/registry.py

@@ -1,14 +0,0 @@
-_model_entrypoints = {}
-
-
-def register_body(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/body/transformer_blocks.py

@@ -1,370 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-# Modified by Bowen Cheng from: https://github.com/facebookresearch/detr/blob/master/models/transformer.py
-"""
-Transformer class.
-
-Copy-paste from torch.nn.Transformer with modifications:
-    * positional encodings are passed in MHattention
-    * extra LN at the end of encoder is removed
-    * decoder returns a stack of activations from all decoding layers
-"""
-import copy
-from typing import List, Optional
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor, nn
-
-
-class Transformer(nn.Module):
-    def __init__(
-        self,
-        d_model=512,
-        nhead=8,
-        num_encoder_layers=6,
-        num_decoder_layers=6,
-        dim_feedforward=2048,
-        dropout=0.1,
-        activation="relu",
-        normalize_before=False,
-        return_intermediate_dec=False,
-    ):
-        super().__init__()
-
-        encoder_layer = TransformerEncoderLayer(
-            d_model, nhead, dim_feedforward, dropout, activation, normalize_before
-        )
-        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
-        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
-
-        decoder_layer = TransformerDecoderLayer(
-            d_model, nhead, dim_feedforward, dropout, activation, normalize_before
-        )
-        decoder_norm = nn.LayerNorm(d_model)
-        self.decoder = TransformerDecoder(
-            decoder_layer,
-            num_decoder_layers,
-            decoder_norm,
-            return_intermediate=return_intermediate_dec,
-        )
-
-        self._reset_parameters()
-
-        self.d_model = d_model
-        self.nhead = nhead
-
-    def _reset_parameters(self):
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-
-    def forward(self, src, mask, query_embed, pos_embed):
-        # flatten NxCxHxW to HWxNxC
-        bs, c, h, w = src.shape
-        src = src.flatten(2).permute(2, 0, 1)
-        pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
-        query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)
-        if mask is not None:
-            mask = mask.flatten(1)
-
-        tgt = torch.zeros_like(query_embed)
-        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
-        hs = self.decoder(
-            tgt, memory, memory_key_padding_mask=mask, pos=pos_embed, query_pos=query_embed
-        )
-        return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w)
-
-
-class TransformerEncoder(nn.Module):
-    def __init__(self, encoder_layer, num_layers, norm=None):
-        super().__init__()
-        self.layers = _get_clones(encoder_layer, num_layers)
-        self.num_layers = num_layers
-        self.norm = norm
-
-    def forward(
-        self,
-        src,
-        mask: Optional[Tensor] = None,
-        src_key_padding_mask: Optional[Tensor] = None,
-        pos: Optional[Tensor] = None,
-    ):
-        output = src
-
-        for layer in self.layers:
-            output = layer(
-                output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos
-            )
-
-        if self.norm is not None:
-            output = self.norm(output)
-
-        return output
-
-
-class TransformerDecoder(nn.Module):
-    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
-        super().__init__()
-        self.layers = _get_clones(decoder_layer, num_layers)
-        self.num_layers = num_layers
-        self.norm = norm
-        self.return_intermediate = return_intermediate
-
-    def forward(
-        self,
-        tgt,
-        memory,
-        tgt_mask: Optional[Tensor] = None,
-        memory_mask: Optional[Tensor] = None,
-        tgt_key_padding_mask: Optional[Tensor] = None,
-        memory_key_padding_mask: Optional[Tensor] = None,
-        pos: Optional[Tensor] = None,
-        query_pos: Optional[Tensor] = None,
-    ):
-        output = tgt
-
-        intermediate = []
-
-        for layer in self.layers:
-            output = layer(
-                output,
-                memory,
-                tgt_mask=tgt_mask,
-                memory_mask=memory_mask,
-                tgt_key_padding_mask=tgt_key_padding_mask,
-                memory_key_padding_mask=memory_key_padding_mask,
-                pos=pos,
-                query_pos=query_pos,
-            )
-            if self.return_intermediate:
-                intermediate.append(self.norm(output))
-
-        if self.norm is not None:
-            output = self.norm(output)
-            if self.return_intermediate:
-                intermediate.pop()
-                intermediate.append(output)
-
-        if self.return_intermediate:
-            return torch.stack(intermediate)
-
-        return output.unsqueeze(0)
-
-
-class TransformerEncoderLayer(nn.Module):
-    def __init__(
-        self,
-        d_model,
-        nhead,
-        dim_feedforward=2048,
-        dropout=0.1,
-        activation="relu",
-        normalize_before=False,
-    ):
-        super().__init__()
-        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
-        # Implementation of Feedforward model
-        self.linear1 = nn.Linear(d_model, dim_feedforward)
-        self.dropout = nn.Dropout(dropout)
-        self.linear2 = nn.Linear(dim_feedforward, d_model)
-
-        self.norm1 = nn.LayerNorm(d_model)
-        self.norm2 = nn.LayerNorm(d_model)
-        self.dropout1 = nn.Dropout(dropout)
-        self.dropout2 = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(
-        self,
-        src,
-        src_mask: Optional[Tensor] = None,
-        src_key_padding_mask: Optional[Tensor] = None,
-        pos: Optional[Tensor] = None,
-    ):
-        q = k = self.with_pos_embed(src, pos)
-
-        src2 = self.self_attn(
-            q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask
-        )[0]
-        src = src + self.dropout1(src2)
-        src = self.norm1(src)
-        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
-        src = src + self.dropout2(src2)
-        src = self.norm2(src)
-        return src
-
-    def forward_pre(
-        self,
-        src,
-        src_mask: Optional[Tensor] = None,
-        src_key_padding_mask: Optional[Tensor] = None,
-        pos: Optional[Tensor] = None,
-    ):
-        src2 = self.norm1(src)
-        q = k = self.with_pos_embed(src2, pos)
-        src2 = self.self_attn(
-            q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask
-        )[0]
-        src = src + self.dropout1(src2)
-        src2 = self.norm2(src)
-        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
-        src = src + self.dropout2(src2)
-        return src
-
-    def forward(
-        self,
-        src,
-        src_mask: Optional[Tensor] = None,
-        src_key_padding_mask: Optional[Tensor] = None,
-        pos: Optional[Tensor] = None,
-    ):
-        if self.normalize_before:
-            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
-        return self.forward_post(src, src_mask, src_key_padding_mask, pos)
-
-
-class TransformerDecoderLayer(nn.Module):
-    def __init__(
-        self,
-        d_model,
-        nhead,
-        dim_feedforward=2048,
-        dropout=0.1,
-        activation="relu",
-        normalize_before=False,
-    ):
-        super().__init__()
-        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
-        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
-        # Implementation of Feedforward model
-        self.linear1 = nn.Linear(d_model, dim_feedforward)
-        self.dropout = nn.Dropout(dropout)
-        self.linear2 = nn.Linear(dim_feedforward, d_model)
-
-        self.norm1 = nn.LayerNorm(d_model)
-        self.norm2 = nn.LayerNorm(d_model)
-        self.norm3 = nn.LayerNorm(d_model)
-        self.dropout1 = nn.Dropout(dropout)
-        self.dropout2 = nn.Dropout(dropout)
-        self.dropout3 = nn.Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-        self.normalize_before = normalize_before
-
-    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
-        return tensor if pos is None else tensor + pos
-
-    def forward_post(
-        self,
-        tgt,
-        memory,
-        tgt_mask: Optional[Tensor] = None,
-        memory_mask: Optional[Tensor] = None,
-        tgt_key_padding_mask: Optional[Tensor] = None,
-        memory_key_padding_mask: Optional[Tensor] = None,
-        pos: Optional[Tensor] = None,
-        query_pos: Optional[Tensor] = None,
-    ):
-        q = k = self.with_pos_embed(tgt, query_pos)
-        tgt2 = self.self_attn(
-            q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
-        )[0]
-        tgt = tgt + self.dropout1(tgt2)
-        tgt = self.norm1(tgt)
-        tgt2 = self.multihead_attn(
-            query=self.with_pos_embed(tgt, query_pos),
-            key=self.with_pos_embed(memory, pos),
-            value=memory,
-            attn_mask=memory_mask,
-            key_padding_mask=memory_key_padding_mask,
-        )[0]
-        tgt = tgt + self.dropout2(tgt2)
-        tgt = self.norm2(tgt)
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
-        tgt = tgt + self.dropout3(tgt2)
-        tgt = self.norm3(tgt)
-        return tgt
-
-    def forward_pre(
-        self,
-        tgt,
-        memory,
-        tgt_mask: Optional[Tensor] = None,
-        memory_mask: Optional[Tensor] = None,
-        tgt_key_padding_mask: Optional[Tensor] = None,
-        memory_key_padding_mask: Optional[Tensor] = None,
-        pos: Optional[Tensor] = None,
-        query_pos: Optional[Tensor] = None,
-    ):
-        tgt2 = self.norm1(tgt)
-        q = k = self.with_pos_embed(tgt2, query_pos)
-        tgt2 = self.self_attn(
-            q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
-        )[0]
-        tgt = tgt + self.dropout1(tgt2)
-        tgt2 = self.norm2(tgt)
-        tgt2 = self.multihead_attn(
-            query=self.with_pos_embed(tgt2, query_pos),
-            key=self.with_pos_embed(memory, pos),
-            value=memory,
-            attn_mask=memory_mask,
-            key_padding_mask=memory_key_padding_mask,
-        )[0]
-        tgt = tgt + self.dropout2(tgt2)
-        tgt2 = self.norm3(tgt)
-        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
-        tgt = tgt + self.dropout3(tgt2)
-        return tgt
-
-    def forward(
-        self,
-        tgt,
-        memory,
-        tgt_mask: Optional[Tensor] = None,
-        memory_mask: Optional[Tensor] = None,
-        tgt_key_padding_mask: Optional[Tensor] = None,
-        memory_key_padding_mask: Optional[Tensor] = None,
-        pos: Optional[Tensor] = None,
-        query_pos: Optional[Tensor] = None,
-    ):
-        if self.normalize_before:
-            return self.forward_pre(
-                tgt,
-                memory,
-                tgt_mask,
-                memory_mask,
-                tgt_key_padding_mask,
-                memory_key_padding_mask,
-                pos,
-                query_pos,
-            )
-        return self.forward_post(
-            tgt,
-            memory,
-            tgt_mask,
-            memory_mask,
-            tgt_key_padding_mask,
-            memory_key_padding_mask,
-            pos,
-            query_pos,
-        )
-
-
-def _get_clones(module, N):
-    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
-
-
-def _get_activation_fn(activation):
-    """Return an activation function given a string"""
-    if activation == "relu":
-        return F.relu
-    if activation == "gelu":
-        return F.gelu
-    if activation == "glu":
-        return F.glu
-    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")

xdecoder/body/xdecoder_head.py

@@ -1,123 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-# --------------------------------------------------------
-# X-Decoder -- Generalized Decoding for Pixel, Image, and Language
-# Copyright (c) 2022 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Jianwei Yang (jianwyan@microsoft.com), Xueyan Zou (xueyan@cs.wisc.edu)
-# --------------------------------------------------------
-
-from typing import Dict
-
-from torch import nn
-
-from detectron2.layers import ShapeSpec
-
-from .registry import register_body
-from .encoder import build_encoder
-from .decoder import build_decoder
-from ..utils import configurable
-
-
-class XDecoderHead(nn.Module):
-
-    @configurable
-    def __init__(
-        self,
-        input_shape: Dict[str, ShapeSpec],
-        *,
-        num_classes: int,
-        pixel_decoder: nn.Module,
-        loss_weight: float = 1.0,
-        ignore_value: int = -1,
-        # extra parameters
-        transformer_predictor: nn.Module,
-        transformer_in_feature: str,
-    ):
-        """
-        NOTE: this interface is experimental.
-        Args:
-            input_shape: shapes (channels and stride) of the input features
-            num_classes: number of classes to predict
-            pixel_decoder: the pixel decoder module
-            loss_weight: loss weight
-            ignore_value: category id to be ignored during training.
-            transformer_predictor: the transformer decoder that makes prediction
-            transformer_in_feature: input feature name to the transformer_predictor
-        """
-        super().__init__()
-
-        input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
-        self.in_features = [k for k, v in input_shape]
-        feature_strides = [v.stride for k, v in input_shape]
-        feature_channels = [v.channels for k, v in input_shape]
-
-        self.ignore_value = ignore_value
-        self.common_stride = 4
-        self.loss_weight = loss_weight
-
-        self.pixel_decoder = pixel_decoder
-        self.predictor = transformer_predictor
-        self.transformer_in_feature = transformer_in_feature
-
-        self.num_classes = num_classes
-
-    @classmethod
-    def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec], lang_encoder: nn.Module, extra: dict):
-
-        in_features_type = cfg['MODEL']['DECODER']['TRANSFORMER_IN_FEATURE']
-        enc_cfg = cfg['MODEL']['ENCODER']
-        dec_cfg = cfg['MODEL']['DECODER']
-
-        # figure out in_channels to transformer predictor
-        if in_features_type == "transformer_encoder":
-            transformer_predictor_in_channels = enc_cfg['CONVS_DIM']
-        elif in_features_type == "pixel_embedding":
-            transformer_predictor_in_channels = enc_cfg['MASK_DIM']
-        elif in_features_type == "multi_scale_pixel_decoder":  # for maskformer2
-            transformer_predictor_in_channels = enc_cfg['CONVS_DIM']
-        else:
-            transformer_predictor_in_channels = input_shape[dec_cfg['TRANSFORMER_IN_FEATURE']].channels
-
-        return {
-            "input_shape": {
-                k: v for k, v in input_shape.items() if k in enc_cfg['IN_FEATURES']
-            },
-            "ignore_value": enc_cfg['IGNORE_VALUE'],
-            "num_classes": enc_cfg.get('NUM_CLASSES', None),
-            "pixel_decoder": build_encoder(cfg, input_shape),
-            "loss_weight": enc_cfg['LOSS_WEIGHT'],
-            "transformer_in_feature": dec_cfg['TRANSFORMER_IN_FEATURE'],
-            "transformer_predictor": build_decoder(
-                cfg,
-                transformer_predictor_in_channels,
-                lang_encoder,
-                mask_classification=True,
-                extra=extra,
-            ),
-        }
-
-    def forward(self, features, mask=None, target_queries=None, target_vlp=None, task='seg', extra={}):
-        return self.layers(features, mask, target_queries, target_vlp, task, extra)
-
-    def layers(self, features, mask=None, target_queries=None, target_vlp=None, task='seg', extra={}):
-        mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
-        
-        if self.transformer_in_feature == "multi_scale_pixel_decoder":
-            predictions = self.predictor(multi_scale_features, mask_features, mask, target_queries, target_vlp, task, extra)
-        else:
-            if self.transformer_in_feature == "transformer_encoder":
-                assert (
-                    transformer_encoder_features is not None
-                ), "Please use the TransformerEncoderPixelDecoder."
-                predictions = self.predictor(transformer_encoder_features, mask_features, mask)
-            elif self.transformer_in_feature == "pixel_embedding":
-                predictions = self.predictor(mask_features, mask_features, mask)
-            else:
-                predictions = self.predictor(features[self.transformer_in_feature], mask_features, mask)
-        return predictions
-
-
-@register_body
-def get_xdecoder_head(cfg, input_shape, lang_encoder, extra):
-    return XDecoderHead(cfg, input_shape, lang_encoder, extra)

xdecoder/language/LangEncoder/__init__.py

@@ -1,8 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from .build import build_lang_encoder
-from .build import build_tokenizer
-
-from .transformer import *

xdecoder/language/LangEncoder/build.py

@@ -1,36 +0,0 @@
-import os
-
-from transformers import CLIPTokenizer, CLIPTokenizerFast
-from transformers import AutoTokenizer
-
-from .registry import lang_encoders
-from .registry import is_lang_encoder
-
-
-def build_lang_encoder(config_encoder, tokenizer, verbose, **kwargs):
-    model_name = config_encoder['NAME']
-
-    if not is_lang_encoder(model_name):
-        raise ValueError(f'Unkown model: {model_name}')
-
-    return lang_encoders(model_name)(config_encoder, tokenizer, verbose, **kwargs)
-
-
-def build_tokenizer(config_encoder):
-    tokenizer = None
-    os.environ['TOKENIZERS_PARALLELISM'] = 'true'
-    if config_encoder['TOKENIZER'] == 'clip':
-        pretrained_tokenizer = config_encoder.get(
-            'PRETRAINED_TOKENIZER', 'openai/clip-vit-base-patch32'
-        )
-        tokenizer = CLIPTokenizer.from_pretrained(pretrained_tokenizer)
-        tokenizer.add_special_tokens({'cls_token': tokenizer.eos_token})
-    elif config_encoder['TOKENIZER'] == 'clip-fast':
-        pretrained_tokenizer = config_encoder.get(
-            'PRETRAINED_TOKENIZER', 'openai/clip-vit-base-patch32'
-        )
-        tokenizer = CLIPTokenizerFast.from_pretrained(pretrained_tokenizer, from_slow=True)
-    else:
-        tokenizer = AutoTokenizer.from_pretrained(config_encoder['TOKENIZER'])
-
-    return tokenizer

xdecoder/language/LangEncoder/registry.py

@@ -1,18 +0,0 @@
-_lang_encoders = {}
-
-
-def register_lang_encoder(fn):
-    module_name_split = fn.__module__.split('.')
-    model_name = module_name_split[-1]
-
-    _lang_encoders[model_name] = fn
-
-    return fn
-
-
-def lang_encoders(model_name):
-    return _lang_encoders[model_name]
-
-
-def is_lang_encoder(model_name):
-    return model_name in _lang_encoders

xdecoder/language/LangEncoder/transformer.py

@@ -1,222 +0,0 @@
-from collections import OrderedDict
-from typing import Tuple, Union
-import logging
-import os
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-from timm.models.layers import DropPath, trunc_normal_
-
-from .registry import register_lang_encoder
-from utils.distributed import is_main_process
-from utils.model import register_norm_module
-
-logger = logging.getLogger(__name__)
-
-
-@register_norm_module
-class LayerNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-12):
-        """Construct a layernorm module in the TF style (epsilon inside the square root).
-        """
-        super(LayerNorm, self).__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.bias = nn.Parameter(torch.zeros(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, x):
-        pdtype = x.dtype
-        x = x.float()
-        u = x.mean(-1, keepdim=True)
-        s = (x - u).pow(2).mean(-1, keepdim=True)
-        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
-        return self.weight * x.to(pdtype) + self.bias
-
-
-class QuickGELU(nn.Module):
-    def forward(self, x: torch.Tensor):
-        return x * torch.sigmoid(1.702 * x)
-
-
-class ResidualAttentionBlock(nn.Module):
-    def __init__(self,
-                 d_model: int,
-                 n_head: int,
-                 attn_mask: torch.Tensor = None,
-                 drop_path: float = 0.0):
-        super().__init__()
-
-        self.attn = nn.MultiheadAttention(d_model, n_head)
-        self.ln_1 = LayerNorm(d_model)
-        self.mlp = nn.Sequential(OrderedDict([
-            ("c_fc", nn.Linear(d_model, d_model * 4)),
-            ("gelu", QuickGELU()),
-            ("c_proj", nn.Linear(d_model * 4, d_model))
-        ]))
-        self.ln_2 = LayerNorm(d_model)
-        self.attn_mask = attn_mask
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-
-    def attention(self, x: torch.Tensor, key_padding_mask: torch.Tensor = None):
-        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) \
-            if self.attn_mask is not None else None
-
-
-        return self.attn(
-            x, x, x,
-            key_padding_mask=key_padding_mask,
-            need_weights=False,
-            attn_mask=self.attn_mask
-        )[0]
-
-    def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor = None):
-        x = x + self.drop_path(self.attention(self.ln_1(x), key_padding_mask=key_padding_mask))
-        x = x + self.drop_path(self.mlp(self.ln_2(x)))
-        return x
-
-
-class Transformer(nn.Module):
-    def __init__(self,
-                 context_length: int,
-                 vocab_size: int,
-                 width: int,
-                 layers: int,
-                 heads: int,
-                 drop_path: float = 0.0,
-                 autogressive: bool =True):
-        super().__init__()
-
-        self.token_embedding = nn.Embedding(vocab_size, width)
-
-        self.context_length = context_length
-        self.positional_embedding = nn.Parameter(
-            torch.empty(self.context_length, width)
-        )
-
-        self.width = width
-        self.layers = layers
-        self.autogressive = autogressive
-        attn_mask = self.build_attention_mask() if autogressive else None
-        dpr = [x.item() for x in torch.linspace(0, drop_path, layers)]  # stochastic depth decay rule
-        self.resblocks = nn.ModuleList(
-            [
-                ResidualAttentionBlock(width, heads, attn_mask, dpr[i])
-                for i in range(layers)
-            ]
-        )
-
-        self.ln_final = LayerNorm(width)
-
-        trunc_normal_(self.positional_embedding, std=.02)
-        # nn.init.normal_(self.token_embedding, std=.02)
-        trunc_normal_(self.token_embedding.weight, std=.02)
-        self.apply(self._init_weights)
-
-    @property
-    def dim_out(self):
-        return self.width
-
-    def build_attention_mask(self):
-        # lazily create causal attention mask, with full attention between the vision tokens
-        # pytorch uses additive attention mask; fill with -inf
-        mask = torch.empty(self.context_length, self.context_length)
-        mask.fill_(float("-inf"))
-        mask.triu_(1)  # zero out the lower diagonal
-        return mask
-
-    def _init_weights(self, m):
-        if isinstance(m, (nn.Linear, nn.Conv2d)):
-            if is_main_process():
-                logger.info('=> init weight of Linear/Conv2d from trunc norm')
-            trunc_normal_(m.weight, std=0.02)
-            if m.bias is not None:
-                if is_main_process():
-                    logger.info('=> init bias of Linear/Conv2d to zeros')
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)):
-            nn.init.constant_(m.bias, 0)
-
-    def load_pretrained(self, pretrained='', pretrained_layers=[], verbose=True):
-        if os.path.isfile(pretrained):
-            pretrained_dict = torch.load(pretrained, map_location='cpu')
-            logging.info(f'=> loading pretrained model {pretrained}')
-            model_dict = self.state_dict()
-            stripped_key = lambda x: x[13:] if x.startswith('lang_encoder.') else x
-            pretrained_dict = {
-                stripped_key(k): v for k, v in pretrained_dict.items()
-                if stripped_key(k) in model_dict.keys()
-            }
-            need_init_state_dict = {}
-            for k, v in pretrained_dict.items():
-                need_init = (
-                    k.split('.')[0] in pretrained_layers
-                    or pretrained_layers[0] == '*'
-                )
-                if need_init:
-                    if verbose:
-                        logger.info(f'=> init {k} from {pretrained}')
-
-                    if 'positional_embedding' in k and v.size() != model_dict[k].size():
-                        positional_embedding_pretrained = v
-                        positional_embedding_current = model_dict[k]
-                        L1, nH1 = positional_embedding_pretrained.size()
-                        L2, nH2 = positional_embedding_current.size()
-                        if nH1 != nH2:
-                            logger.info(f"Error in loading {k}, passing")
-                        else:
-                            if L1 != L2:
-                                logger.info(
-                                    '=> load_pretrained: resized variant: {} to {}'
-                                        .format((L1, nH1), (L2, nH2))
-                                )
-
-                                posemb = positional_embedding_pretrained.float()
-                                posemb_grid = posemb.unsqueeze(dim=0).permute(0, 2, 1)
-                                posemb_grid = torch.nn.functional.interpolate(posemb_grid, size=L2, mode='linear')
-                                posemb_grid = posemb_grid.permute(0, 2, 1).squeeze(dim=0)
-                                v = posemb_grid
-
-                    need_init_state_dict[k] = v
-
-            self.load_state_dict(need_init_state_dict, strict=False)
-
-
-    @torch.jit.ignore
-    def no_weight_decay(self):
-        return {
-            'positional_embedding',
-            'token_embedding',
-        }
-
-    def forward(self, input_ids, attention_mask=None):
-        key_padding_mask = (attention_mask == 0) if (not self.autogressive and attention_mask is not None) else None
-        # key_padding_mask = (input_ids == 0) if not self.autogressive else None
-        x = self.token_embedding(input_ids)  # [batch_size, n_ctx, d_model]
-        x = x + self.positional_embedding
-        x = x.permute(1, 0, 2)  # NLD -> LND
-        for block in self.resblocks:
-            x = block(x, key_padding_mask)
-        x = x.permute(1, 0, 2)  # LND -> NLD
-
-        x = self.ln_final(x)
-
-        return {'last_hidden_state': x}
-
-
-@register_lang_encoder
-def lang_encoder(config_encoder, tokenizer, verbose, **kwargs):
-    transformer = Transformer(
-        context_length=config_encoder['CONTEXT_LENGTH'],
-        vocab_size=tokenizer.vocab_size,
-        width=config_encoder['WIDTH'],
-        layers=config_encoder['LAYERS'],
-        heads=config_encoder['HEADS'],
-        autogressive=config_encoder.get('AUTOGRESSIVE', True)
-    )
-
-    if config_encoder.get('LOAD_PRETRAINED', False):
-        transformer.load_pretrained(config_encoder['PRETRAINED'], config_encoder.get('PRETRAINED_LAYERS', ['*']))
-    return transformer

xdecoder/language/__init__.py

@@ -1,3 +0,0 @@
-from .fixvlpencoder import *
-from .vlpencoder import *
-from .build import build_language_encoder

xdecoder/language/build.py

@@ -1,11 +0,0 @@
-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

@@ -1,35 +0,0 @@
-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

@@ -1,225 +0,0 @@
-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

@@ -1,64 +0,0 @@
-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

@@ -1,13 +0,0 @@
-_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

@@ -1,168 +0,0 @@
-
-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)

xdecoder/modules/__init__.py

@@ -1,3 +0,0 @@
-from .position_encoding import *
-from .attention import *
-from .postprocessing import *

xdecoder/modules/attention.py

@@ -1,489 +0,0 @@
-# Code copy from PyTorch, modified by Xueyan Zou
-
-import warnings
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-from torch import Tensor
-from torch.nn.init import constant_, xavier_normal_, xavier_uniform_
-from torch.nn.parameter import Parameter
-from torch.overrides import has_torch_function, handle_torch_function
-from torch.nn.functional import pad, linear, softmax, dropout
-
-
-def multi_head_attention_forward(
-    query: Tensor,
-    key: Tensor,
-    value: Tensor,
-    embed_dim_to_check: int,
-    num_heads: int,
-    in_proj_weight: Tensor,
-    in_proj_bias: Tensor,
-    bias_k: Optional[Tensor],
-    bias_v: Optional[Tensor],
-    add_zero_attn: bool,
-    dropout_p: float,
-    out_proj_weight: Tensor,
-    out_proj_bias: Tensor,
-    training: bool = True,
-    key_padding_mask: Optional[Tensor] = None,
-    need_weights: bool = True,
-    attn_mask: Optional[Tensor] = None,
-    use_separate_proj_weight: bool = False,
-    q_proj_weight: Optional[Tensor] = None,
-    k_proj_weight: Optional[Tensor] = None,
-    v_proj_weight: Optional[Tensor] = None,
-    static_k: Optional[Tensor] = None,
-    static_v: Optional[Tensor] = None,
-) -> Tuple[Tensor, Optional[Tensor]]:
-    r"""
-    Args:
-        query, key, value: map a query and a set of key-value pairs to an output.
-            See "Attention Is All You Need" for more details.
-        embed_dim_to_check: total dimension of the model.
-        num_heads: parallel attention heads.
-        in_proj_weight, in_proj_bias: input projection weight and bias.
-        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
-        add_zero_attn: add a new batch of zeros to the key and
-                       value sequences at dim=1.
-        dropout_p: probability of an element to be zeroed.
-        out_proj_weight, out_proj_bias: the output projection weight and bias.
-        training: apply dropout if is ``True``.
-        key_padding_mask: if provided, specified padding elements in the key will
-            be ignored by the attention. This is an binary mask. When the value is True,
-            the corresponding value on the attention layer will be filled with -inf.
-        need_weights: output attn_output_weights.
-        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
-            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
-        use_separate_proj_weight: the function accept the proj. weights for query, key,
-            and value in different forms. If false, in_proj_weight will be used, which is
-            a combination of q_proj_weight, k_proj_weight, v_proj_weight.
-        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
-        static_k, static_v: static key and value used for attention operators.
-
-
-    Shape:
-        Inputs:
-        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
-          If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
-          will be unchanged. If a BoolTensor is provided, the positions with the
-          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
-        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
-          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
-          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
-          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
-          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
-          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
-          is provided, it will be added to the attention weight.
-        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
-          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
-        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
-          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
-
-        Outputs:
-        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
-          E is the embedding dimension.
-        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
-          L is the target sequence length, S is the source sequence length.
-    """
-    tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v, out_proj_weight, out_proj_bias)
-    if has_torch_function(tens_ops):
-        return handle_torch_function(
-            multi_head_attention_forward,
-            tens_ops,
-            query,
-            key,
-            value,
-            embed_dim_to_check,
-            num_heads,
-            in_proj_weight,
-            in_proj_bias,
-            bias_k,
-            bias_v,
-            add_zero_attn,
-            dropout_p,
-            out_proj_weight,
-            out_proj_bias,
-            training=training,
-            key_padding_mask=key_padding_mask,
-            need_weights=need_weights,
-            attn_mask=attn_mask,
-            use_separate_proj_weight=use_separate_proj_weight,
-            q_proj_weight=q_proj_weight,
-            k_proj_weight=k_proj_weight,
-            v_proj_weight=v_proj_weight,
-            static_k=static_k,
-            static_v=static_v,
-        )
-    tgt_len, bsz, embed_dim = query.size()
-    assert embed_dim == embed_dim_to_check
-    # allow MHA to have different sizes for the feature dimension
-    assert key.size(0) == value.size(0) and key.size(1) == value.size(1)
-
-    head_dim = embed_dim // num_heads
-    assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
-    scaling = float(head_dim) ** -0.5
-
-    if not use_separate_proj_weight:
-        if (query is key or torch.equal(query, key)) and (key is value or torch.equal(key, value)):
-            # self-attention
-            q, k, v = linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)
-
-        elif key is value or torch.equal(key, value):
-            # encoder-decoder attention
-            # This is inline in_proj function with in_proj_weight and in_proj_bias
-            _b = in_proj_bias
-            _start = 0
-            _end = embed_dim
-            _w = in_proj_weight[_start:_end, :]
-            if _b is not None:
-                _b = _b[_start:_end]
-            q = linear(query, _w, _b)
-
-            if key is None:
-                assert value is None
-                k = None
-                v = None
-            else:
-
-                # This is inline in_proj function with in_proj_weight and in_proj_bias
-                _b = in_proj_bias
-                _start = embed_dim
-                _end = None
-                _w = in_proj_weight[_start:, :]
-                if _b is not None:
-                    _b = _b[_start:]
-                k, v = linear(key, _w, _b).chunk(2, dim=-1)
-
-        else:
-            # This is inline in_proj function with in_proj_weight and in_proj_bias
-            _b = in_proj_bias
-            _start = 0
-            _end = embed_dim
-            _w = in_proj_weight[_start:_end, :]
-            if _b is not None:
-                _b = _b[_start:_end]
-            q = linear(query, _w, _b)
-
-            # This is inline in_proj function with in_proj_weight and in_proj_bias
-            _b = in_proj_bias
-            _start = embed_dim
-            _end = embed_dim * 2
-            _w = in_proj_weight[_start:_end, :]
-            if _b is not None:
-                _b = _b[_start:_end]
-            k = linear(key, _w, _b)
-
-            # This is inline in_proj function with in_proj_weight and in_proj_bias
-            _b = in_proj_bias
-            _start = embed_dim * 2
-            _end = None
-            _w = in_proj_weight[_start:, :]
-            if _b is not None:
-                _b = _b[_start:]
-            v = linear(value, _w, _b)
-    else:
-        q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight)
-        len1, len2 = q_proj_weight_non_opt.size()
-        assert len1 == embed_dim and len2 == query.size(-1)
-
-        k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight)
-        len1, len2 = k_proj_weight_non_opt.size()
-        assert len1 == embed_dim and len2 == key.size(-1)
-
-        v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight)
-        len1, len2 = v_proj_weight_non_opt.size()
-        assert len1 == embed_dim and len2 == value.size(-1)
-
-        if in_proj_bias is not None:
-            q = linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim])
-            k = linear(key, k_proj_weight_non_opt, in_proj_bias[embed_dim : (embed_dim * 2)])
-            v = linear(value, v_proj_weight_non_opt, in_proj_bias[(embed_dim * 2) :])
-        else:
-            q = linear(query, q_proj_weight_non_opt, in_proj_bias)
-            k = linear(key, k_proj_weight_non_opt, in_proj_bias)
-            v = linear(value, v_proj_weight_non_opt, in_proj_bias)
-    q = q * scaling
-
-    if attn_mask is not None:
-        assert (
-            attn_mask.dtype == torch.float32
-            or attn_mask.dtype == torch.float64
-            or attn_mask.dtype == torch.float16
-            or attn_mask.dtype == torch.uint8
-            or attn_mask.dtype == torch.bool
-        ), "Only float, byte, and bool types are supported for attn_mask, not {}".format(attn_mask.dtype)
-        if attn_mask.dtype == torch.uint8:
-            warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
-            attn_mask = attn_mask.to(torch.bool)
-
-        if attn_mask.dim() == 2:
-            attn_mask = attn_mask.unsqueeze(0)
-            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
-                raise RuntimeError("The size of the 2D attn_mask is not correct.")
-        elif attn_mask.dim() == 3:
-            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
-                raise RuntimeError("The size of the 3D attn_mask is not correct.")
-        else:
-            raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))
-        # attn_mask's dim is 3 now.
-
-    # convert ByteTensor key_padding_mask to bool
-    if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
-        warnings.warn(
-            "Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead."
-        )
-        key_padding_mask = key_padding_mask.to(torch.bool)
-
-    if bias_k is not None and bias_v is not None:
-        if static_k is None and static_v is None:
-            k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
-            v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
-            if attn_mask is not None:
-                attn_mask = pad(attn_mask, (0, 1))
-            if key_padding_mask is not None:
-                key_padding_mask = pad(key_padding_mask, (0, 1))
-        else:
-            assert static_k is None, "bias cannot be added to static key."
-            assert static_v is None, "bias cannot be added to static value."
-    else:
-        assert bias_k is None
-        assert bias_v is None
-
-    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
-    if k is not None:
-        k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
-    if v is not None:
-        v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
-
-    if static_k is not None:
-        assert static_k.size(0) == bsz * num_heads
-        assert static_k.size(2) == head_dim
-        k = static_k
-
-    if static_v is not None:
-        assert static_v.size(0) == bsz * num_heads
-        assert static_v.size(2) == head_dim
-        v = static_v
-
-    src_len = k.size(1)
-
-    if key_padding_mask is not None:
-        # assert key_padding_mask.size(0) == bsz
-        assert key_padding_mask.size(1) == src_len
-
-    if add_zero_attn:
-        src_len += 1
-        k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1)
-        v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1)
-        if attn_mask is not None:
-            attn_mask = pad(attn_mask, (0, 1))
-        if key_padding_mask is not None:
-            key_padding_mask = pad(key_padding_mask, (0, 1))
-
-    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
-    assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
-
-    if attn_mask is not None:
-        if attn_mask.dtype == torch.bool:
-            attn_output_weights.masked_fill_(attn_mask, float("-inf"))
-        else:
-            attn_output_weights += attn_mask
-
-    if key_padding_mask is not None:
-        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
-        attn_output_weights = attn_output_weights.masked_fill(
-            key_padding_mask.unsqueeze(1),
-            float("-inf"),
-        )
-        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)
-
-    attn_output_weights = softmax(attn_output_weights, dim=-1)
-    attn_output_weights = dropout(attn_output_weights, p=dropout_p, training=training)
-
-    attn_output = torch.bmm(attn_output_weights, v)
-    assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]
-    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
-    attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
-
-    if need_weights:
-        # average attention weights over heads
-        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
-        return attn_output, attn_output_weights.sum(dim=1) / num_heads
-    else:
-        return attn_output, None
-
-
-# This class exists solely for Transformer; it has an annotation stating
-# that bias is never None, which appeases TorchScript
-class _LinearWithBias(nn.Linear):
-    bias: Tensor  # type: ignore
-
-    def __init__(self, in_features: int, out_features: int) -> None:
-        super().__init__(in_features, out_features, bias=True)  # type: ignore
-
-
-class MultiheadAttention(nn.Module):
-    r"""Allows the model to jointly attend to information
-    from different representation subspaces.
-    See `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_
-
-    .. math::
-        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
-
-    where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
-
-    Args:
-        embed_dim: total dimension of the model.
-        num_heads: parallel attention heads.
-        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
-        bias: add bias as module parameter. Default: True.
-        add_bias_kv: add bias to the key and value sequences at dim=0.
-        add_zero_attn: add a new batch of zeros to the key and
-                       value sequences at dim=1.
-        kdim: total number of features in key. Default: None.
-        vdim: total number of features in value. Default: None.
-
-    Note that if :attr:`kdim` and :attr:`vdim` are None, they will be set
-    to :attr:`embed_dim` such that query, key, and value have the same
-    number of features.
-
-    Examples::
-
-        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
-        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
-    """
-    bias_k: Optional[torch.Tensor]
-    bias_v: Optional[torch.Tensor]
-
-    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
-        super(MultiheadAttention, self).__init__()
-        self.embed_dim = embed_dim
-        self.kdim = kdim if kdim is not None else embed_dim
-        self.vdim = vdim if vdim is not None else embed_dim
-        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
-
-        self.num_heads = num_heads
-        self.dropout = dropout
-        self.head_dim = embed_dim // num_heads
-        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
-
-        if self._qkv_same_embed_dim is False:
-            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
-            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
-            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
-            self.register_parameter('in_proj_weight', None)
-        else:
-            self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
-            self.register_parameter('q_proj_weight', None)
-            self.register_parameter('k_proj_weight', None)
-            self.register_parameter('v_proj_weight', None)
-
-        if bias:
-            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
-        else:
-            self.register_parameter('in_proj_bias', None)
-        self.out_proj = _LinearWithBias(embed_dim, embed_dim)
-
-        if add_bias_kv:
-            self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
-            self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
-        else:
-            self.bias_k = self.bias_v = None
-
-        self.add_zero_attn = add_zero_attn
-
-        self._reset_parameters()
-
-    def _reset_parameters(self):
-        if self._qkv_same_embed_dim:
-            xavier_uniform_(self.in_proj_weight)
-        else:
-            xavier_uniform_(self.q_proj_weight)
-            xavier_uniform_(self.k_proj_weight)
-            xavier_uniform_(self.v_proj_weight)
-
-        if self.in_proj_bias is not None:
-            constant_(self.in_proj_bias, 0.)
-            constant_(self.out_proj.bias, 0.)
-        if self.bias_k is not None:
-            xavier_normal_(self.bias_k)
-        if self.bias_v is not None:
-            xavier_normal_(self.bias_v)
-
-    def __setstate__(self, state):
-        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
-        if '_qkv_same_embed_dim' not in state:
-            state['_qkv_same_embed_dim'] = True
-
-        super(MultiheadAttention, self).__setstate__(state)
-
-    def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None,
-                need_weights: bool = True, attn_mask: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:
-        r"""
-    Args:
-        query, key, value: map a query and a set of key-value pairs to an output.
-            See "Attention Is All You Need" for more details.
-        key_padding_mask: if provided, specified padding elements in the key will
-            be ignored by the attention. When given a binary mask and a value is True,
-            the corresponding value on the attention layer will be ignored. When given
-            a byte mask and a value is non-zero, the corresponding value on the attention
-            layer will be ignored
-        need_weights: output attn_output_weights.
-        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
-            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
-
-    Shapes for inputs:
-        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
-          the embedding dimension.
-        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
-          If a ByteTensor is provided, the non-zero positions will be ignored while the position
-          with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
-          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
-        - attn_mask: if a 2D mask: :math:`(L, S)` where L is the target sequence length, S is the
-          source sequence length.
-
-          If a 3D mask: :math:`(N\cdot\text{num\_heads}, L, S)` where N is the batch size, L is the target sequence
-          length, S is the source sequence length. ``attn_mask`` ensure that position i is allowed to attend
-          the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
-          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
-          is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
-          is provided, it will be added to the attention weight.
-
-    Shapes for outputs:
-        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
-          E is the embedding dimension.
-        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
-          L is the target sequence length, S is the source sequence length.
-        """
-        if not self._qkv_same_embed_dim:
-            return multi_head_attention_forward(
-                query, key, value, self.embed_dim, self.num_heads,
-                self.in_proj_weight, self.in_proj_bias,
-                self.bias_k, self.bias_v, self.add_zero_attn,
-                self.dropout, self.out_proj.weight, self.out_proj.bias,
-                training=self.training,
-                key_padding_mask=key_padding_mask, need_weights=need_weights,
-                attn_mask=attn_mask, use_separate_proj_weight=True,
-                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
-                v_proj_weight=self.v_proj_weight)
-        else:
-            return multi_head_attention_forward(
-                query, key, value, self.embed_dim, self.num_heads,
-                self.in_proj_weight, self.in_proj_bias,
-                self.bias_k, self.bias_v, self.add_zero_attn,
-                self.dropout, self.out_proj.weight, self.out_proj.bias,
-                training=self.training,
-                key_padding_mask=key_padding_mask, need_weights=need_weights,
-                attn_mask=attn_mask)

xdecoder/modules/position_encoding.py

@@ -1,64 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-## Modified by Bowen Cheng from: https://github.com/facebookresearch/detr/blob/master/models/position_encoding.py
-"""
-Various positional encodings for the transformer.
-"""
-import math
-
-import torch
-from torch import nn
-
-
-class PositionEmbeddingSine(nn.Module):
-    """
-    This is a more standard version of the position embedding, very similar to the one
-    used by the Attention is all you need paper, generalized to work on images.
-    """
-
-    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
-        super().__init__()
-        self.num_pos_feats = num_pos_feats
-        self.temperature = temperature
-        self.normalize = normalize
-        if scale is not None and normalize is False:
-            raise ValueError("normalize should be True if scale is passed")
-        if scale is None:
-            scale = 2 * math.pi
-        self.scale = scale
-
-    def forward(self, x, mask=None):
-        if mask is None:
-            mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)
-        not_mask = ~mask
-        y_embed = not_mask.cumsum(1, dtype=x.dtype)
-        x_embed = not_mask.cumsum(2, dtype=x.dtype)
-        if self.normalize:
-            eps = 1e-6
-            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
-            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
-
-        dim_t = torch.arange(self.num_pos_feats, dtype=x.dtype, device=x.device)
-        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
-
-        pos_x = x_embed[:, :, :, None] / dim_t
-        pos_y = y_embed[:, :, :, None] / dim_t
-        pos_x = torch.stack(
-            (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4
-        ).flatten(3)
-        pos_y = torch.stack(
-            (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4
-        ).flatten(3)
-        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
-        return pos
-    
-    def __repr__(self, _repr_indent=4):
-        head = "Positional encoding " + self.__class__.__name__
-        body = [
-            "num_pos_feats: {}".format(self.num_pos_feats),
-            "temperature: {}".format(self.temperature),
-            "normalize: {}".format(self.normalize),
-            "scale: {}".format(self.scale),
-        ]
-        # _repr_indent = 4
-        lines = [head] + [" " * _repr_indent + line for line in body]
-        return "\n".join(lines)

xdecoder/modules/postprocessing.py

@@ -1,122 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-import torch
-from torch.nn import functional as F
-
-from detectron2.structures import Instances, ROIMasks
-
-
-# perhaps should rename to "resize_instance"
-def detector_postprocess(
-    results: Instances, output_height: int, output_width: int, mask_threshold: float = 0.5
-):
-    """
-    Resize the output instances.
-    The input images are often resized when entering an object detector.
-    As a result, we often need the outputs of the detector in a different
-    resolution from its inputs.
-
-    This function will resize the raw outputs of an R-CNN detector
-    to produce outputs according to the desired output resolution.
-
-    Args:
-        results (Instances): the raw outputs from the detector.
-            `results.image_size` contains the input image resolution the detector sees.
-            This object might be modified in-place.
-        output_height, output_width: the desired output resolution.
-
-    Returns:
-        Instances: the resized output from the model, based on the output resolution
-    """
-    if isinstance(output_width, torch.Tensor):
-        # This shape might (but not necessarily) be tensors during tracing.
-        # Converts integer tensors to float temporaries to ensure true
-        # division is performed when computing scale_x and scale_y.
-        output_width_tmp = output_width.float()
-        output_height_tmp = output_height.float()
-        new_size = torch.stack([output_height, output_width])
-    else:
-        new_size = (output_height, output_width)
-        output_width_tmp = output_width
-        output_height_tmp = output_height
-
-    scale_x, scale_y = (
-        output_width_tmp / results.image_size[1],
-        output_height_tmp / results.image_size[0],
-    )
-    results = Instances(new_size, **results.get_fields())
-
-    if results.has("pred_boxes"):
-        output_boxes = results.pred_boxes
-    elif results.has("proposal_boxes"):
-        output_boxes = results.proposal_boxes
-    else:
-        output_boxes = None
-    assert output_boxes is not None, "Predictions must contain boxes!"
-
-    output_boxes.scale(scale_x, scale_y)
-    output_boxes.clip(results.image_size)
-
-    results = results[output_boxes.nonempty()]
-
-    if results.has("pred_masks"):
-        if isinstance(results.pred_masks, ROIMasks):
-            roi_masks = results.pred_masks
-        else:
-            # pred_masks is a tensor of shape (N, 1, M, M)
-            roi_masks = ROIMasks(results.pred_masks[:, 0, :, :])
-        results.pred_masks = roi_masks.to_bitmasks(
-            results.pred_boxes, output_height, output_width, mask_threshold
-        ).tensor  # TODO return ROIMasks/BitMask object in the future
-
-    if results.has("pred_keypoints"):
-        results.pred_keypoints[:, :, 0] *= scale_x
-        results.pred_keypoints[:, :, 1] *= scale_y
-
-    return results
-
-def bbox_postprocess(result, input_size, img_size, output_height, output_width):
-    """
-    result: [xc,yc,w,h] range [0,1] to [x1,y1,x2,y2] range [0,w], [0,h]
-    """
-    if result is None:
-        return None
-    
-    scale = torch.tensor([input_size[1], input_size[0], input_size[1], input_size[0]])[None,:].to(result.device)
-    result = result.sigmoid() * scale
-    x1,y1,x2,y2 = result[:,0] - result[:,2]/2, result[:,1] - result[:,3]/2, result[:,0] + result[:,2]/2, result[:,1] + result[:,3]/2
-    h,w = img_size
-
-    x1 = x1.clamp(min=0, max=w)
-    y1 = y1.clamp(min=0, max=h)
-    x2 = x2.clamp(min=0, max=w)
-    y2 = y2.clamp(min=0, max=h)
-
-    box = torch.stack([x1,y1,x2,y2]).permute(1,0)
-    scale = torch.tensor([output_width/w, output_height/h, output_width/w, output_height/h])[None,:].to(result.device)
-    box = box*scale
-    return box
-
-def sem_seg_postprocess(result, img_size, output_height, output_width):
-    """
-    Return semantic segmentation predictions in the original resolution.
-
-    The input images are often resized when entering semantic segmentor. Moreover, in same
-    cases, they also padded inside segmentor to be divisible by maximum network stride.
-    As a result, we often need the predictions of the segmentor in a different
-    resolution from its inputs.
-
-    Args:
-        result (Tensor): semantic segmentation prediction logits. A tensor of shape (C, H, W),
-            where C is the number of classes, and H, W are the height and width of the prediction.
-        img_size (tuple): image size that segmentor is taking as input.
-        output_height, output_width: the desired output resolution.
-
-    Returns:
-        semantic segmentation prediction (Tensor): A tensor of the shape
-            (C, output_height, output_width) that contains per-pixel soft predictions.
-    """
-    result = result[:, : img_size[0], : img_size[1]].expand(1, -1, -1, -1)
-    result = F.interpolate(
-        result, size=(output_height, output_width), mode="bilinear", align_corners=False
-    )[0]
-    return result

xdecoder/utils/__init__.py

@@ -1,4 +0,0 @@
-from .config import *
-from .misc import *
-from .box_ops import *
-from .it_contrastive import *

xdecoder/utils/box_ops.py

@@ -1,93 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-"""
-Utilities for bounding box manipulation and GIoU.
-"""
-import torch
-from torchvision.ops.boxes import box_area
-
-
-def box_cxcywh_to_xyxy(x):
-    x_c, y_c, w, h = x.unbind(-1)
-    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
-         (x_c + 0.5 * w), (y_c + 0.5 * h)]
-    return torch.stack(b, dim=-1)
-
-
-def box_xyxy_to_cxcywh(x):
-    x0, y0, x1, y1 = x.unbind(-1)
-    b = [(x0 + x1) / 2, (y0 + y1) / 2,
-         (x1 - x0), (y1 - y0)]
-    return torch.stack(b, dim=-1)
-
-def box_xywh_to_xyxy(x):
-    x0, y0, x1, y1 = x.unbind(-1)
-    b = [x0, y0, (x0 + x1), (y0 + y1)]
-    return torch.stack(b, dim=-1)
-
-
-# modified from torchvision to also return the union
-def box_iou(boxes1, boxes2):
-    area1 = box_area(boxes1)
-    area2 = box_area(boxes2)
-
-    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
-    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]
-
-    wh = (rb - lt).clamp(min=0)  # [N,M,2]
-    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
-
-    union = area1[:, None] + area2 - inter
-
-    iou = inter / union
-    return iou, union
-
-
-def generalized_box_iou(boxes1, boxes2):
-    """
-    Generalized IoU from https://giou.stanford.edu/
-
-    The boxes should be in [x0, y0, x1, y1] format
-
-    Returns a [N, M] pairwise matrix, where N = len(boxes1)
-    and M = len(boxes2)
-    """
-    # degenerate boxes gives inf / nan results
-    # so do an early check
-    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
-    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
-    iou, union = box_iou(boxes1, boxes2)
-
-    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
-    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
-
-    wh = (rb - lt).clamp(min=0)  # [N,M,2]
-    area = wh[:, :, 0] * wh[:, :, 1]
-
-    return iou - (area - union) / area
-
-
-def masks_to_boxes(masks):
-    """Compute the bounding boxes around the provided masks
-
-    The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.
-
-    Returns a [N, 4] tensors, with the boxes in xyxy format
-    """
-    if masks.numel() == 0:
-        return torch.zeros((0, 4), device=masks.device)
-
-    h, w = masks.shape[-2:]
-
-    y = torch.arange(0, h, dtype=torch.float)
-    x = torch.arange(0, w, dtype=torch.float)
-    y, x = torch.meshgrid(y, x)
-
-    x_mask = (masks * x.unsqueeze(0))
-    x_max = x_mask.flatten(1).max(-1)[0]
-    x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
-
-    y_mask = (masks * y.unsqueeze(0))
-    y_max = y_mask.flatten(1).max(-1)[0]
-    y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
-
-    return torch.stack([x_min, y_min, x_max, y_max], 1)

xdecoder/utils/config.py

@@ -1,140 +0,0 @@
-# -*- coding: utf-8 -*-
-# Copyright (c) Facebook, Inc. and its affiliates.
-
-import functools
-import inspect
-
-def configurable(init_func=None, *, from_config=None):
-    """
-    Decorate a function or a class's __init__ method so that it can be called
-    with a :class:`CfgNode` object using a :func:`from_config` function that translates
-    :class:`CfgNode` to arguments.
-
-    Examples:
-    ::
-        # Usage 1: Decorator on __init__:
-        class A:
-            @configurable
-            def __init__(self, a, b=2, c=3):
-                pass
-
-            @classmethod
-            def from_config(cls, cfg):   # 'cfg' must be the first argument
-                # Returns kwargs to be passed to __init__
-                return {"a": cfg.A, "b": cfg.B}
-
-        a1 = A(a=1, b=2)  # regular construction
-        a2 = A(cfg)       # construct with a cfg
-        a3 = A(cfg, b=3, c=4)  # construct with extra overwrite
-
-        # Usage 2: Decorator on any function. Needs an extra from_config argument:
-        @configurable(from_config=lambda cfg: {"a: cfg.A, "b": cfg.B})
-        def a_func(a, b=2, c=3):
-            pass
-
-        a1 = a_func(a=1, b=2)  # regular call
-        a2 = a_func(cfg)       # call with a cfg
-        a3 = a_func(cfg, b=3, c=4)  # call with extra overwrite
-
-    Args:
-        init_func (callable): a class's ``__init__`` method in usage 1. The
-            class must have a ``from_config`` classmethod which takes `cfg` as
-            the first argument.
-        from_config (callable): the from_config function in usage 2. It must take `cfg`
-            as its first argument.
-    """
-
-    if init_func is not None:
-        assert (
-            inspect.isfunction(init_func)
-            and from_config is None
-            and init_func.__name__ == "__init__"
-        ), "Incorrect use of @configurable. Check API documentation for examples."
-
-        @functools.wraps(init_func)
-        def wrapped(self, *args, **kwargs):
-            try:
-                from_config_func = type(self).from_config
-            except AttributeError as e:
-                raise AttributeError(
-                    "Class with @configurable must have a 'from_config' classmethod."
-                ) from e
-            if not inspect.ismethod(from_config_func):
-                raise TypeError("Class with @configurable must have a 'from_config' classmethod.")
-
-            if _called_with_cfg(*args, **kwargs):
-                explicit_args = _get_args_from_config(from_config_func, *args, **kwargs)
-                init_func(self, **explicit_args)
-            else:
-                init_func(self, *args, **kwargs)
-
-        return wrapped
-
-    else:
-        if from_config is None:
-            return configurable  # @configurable() is made equivalent to @configurable
-        assert inspect.isfunction(
-            from_config
-        ), "from_config argument of configurable must be a function!"
-
-        def wrapper(orig_func):
-            @functools.wraps(orig_func)
-            def wrapped(*args, **kwargs):
-                if _called_with_cfg(*args, **kwargs):
-                    explicit_args = _get_args_from_config(from_config, *args, **kwargs)
-                    return orig_func(**explicit_args)
-                else:
-                    return orig_func(*args, **kwargs)
-
-            wrapped.from_config = from_config
-            return wrapped
-
-        return wrapper
-
-def _called_with_cfg(*args, **kwargs):
-    """
-    Returns:
-        bool: whether the arguments contain CfgNode and should be considered
-            forwarded to from_config.
-    """
-    from omegaconf import DictConfig
-
-    if len(args) and isinstance(args[0], (dict)):
-        return True
-    if isinstance(kwargs.pop("cfg", None), (dict)):
-        return True
-    # `from_config`'s first argument is forced to be "cfg".
-    # So the above check covers all cases.
-    return False
-
-def _get_args_from_config(from_config_func, *args, **kwargs):
-    """
-    Use `from_config` to obtain explicit arguments.
-
-    Returns:
-        dict: arguments to be used for cls.__init__
-    """
-    signature = inspect.signature(from_config_func)
-    if list(signature.parameters.keys())[0] != "cfg":
-        if inspect.isfunction(from_config_func):
-            name = from_config_func.__name__
-        else:
-            name = f"{from_config_func.__self__}.from_config"
-        raise TypeError(f"{name} must take 'cfg' as the first argument!")
-    support_var_arg = any(
-        param.kind in [param.VAR_POSITIONAL, param.VAR_KEYWORD]
-        for param in signature.parameters.values()
-    )
-    if support_var_arg:  # forward all arguments to from_config, if from_config accepts them
-        ret = from_config_func(*args, **kwargs)
-    else:
-        # forward supported arguments to from_config
-        supported_arg_names = set(signature.parameters.keys())
-        extra_kwargs = {}
-        for name in list(kwargs.keys()):
-            if name not in supported_arg_names:
-                extra_kwargs[name] = kwargs.pop(name)
-        ret = from_config_func(*args, **kwargs)
-        # forward the other arguments to __init__
-        ret.update(extra_kwargs)
-    return ret

xdecoder/utils/it_contrastive.py

@@ -1,59 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-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 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
-
-@torch.no_grad()
-def all_gather_nograd(tensor):
-    # from albef
-    """
-    Performs all_gather operation on the provided tensors.
-    *** Warning ***: torch.distributed.all_gather has no gradient.
-    """
-    if get_world_size() > 1:
-        tensors_gather = [torch.ones_like(tensor)
-            for _ in range(torch.distributed.get_world_size())]
-        torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
-
-        tensor = torch.cat(tensors_gather, dim=0)
-    return tensor
-
-def image_text_contrastive_loss(image_feat, text_feat, temperature, image_id=None, text_id=None):
-    # 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())
-    logits /= temperature
-    
-    if image_id is None and text_id is None:
-        gt = torch.arange(logits.shape[0], device=logits.device)
-        loss1 = F.cross_entropy(logits, gt)
-        loss2 = F.cross_entropy(logits.t(), gt)        
-    else:
-        image_id = all_gather_grad(image_id)
-        text_id = all_gather_grad(text_id)
-
-        gt_image = image_id.reshape((-1, 1)) == image_id.reshape((1, -1))
-        gt_text = text_id.reshape((-1, 1)) == text_id.reshape((1, -1))
-        gt = torch.logical_or(gt_image, gt_text)
-
-        loss1 = -torch.sum(gt * F.log_softmax(logits, dim=1)) / gt.sum()
-        loss2 = -torch.sum(gt.t() * F.log_softmax(logits.t(), dim=1)) / gt.sum()
-
-    return (loss1 + loss2) / 2 * get_world_size() # scale it up by the number of GPUs

xdecoder/utils/misc.py

@@ -1,157 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/util/misc.py
-# Modified by Xueyan Zou
-"""
-Misc functions, including distributed helpers.
-
-Mostly copy-paste from torchvision references.
-"""
-from typing import List, Optional
-
-import torch
-import torch.distributed as dist
-import torchvision
-from torch import Tensor
-
-def _max_by_axis(the_list):
-    # type: (List[List[int]]) -> List[int]
-    maxes = the_list[0]
-    for sublist in the_list[1:]:
-        for index, item in enumerate(sublist):
-            maxes[index] = max(maxes[index], item)
-    return maxes
-
-class NestedTensor(object):
-    def __init__(self, tensors, mask: Optional[Tensor]):
-        self.tensors = tensors
-        self.mask = mask
-
-    def to(self, device):
-        # type: (Device) -> NestedTensor # noqa
-        cast_tensor = self.tensors.to(device)
-        mask = self.mask
-        if mask is not None:
-            assert mask is not None
-            cast_mask = mask.to(device)
-        else:
-            cast_mask = None
-        return NestedTensor(cast_tensor, cast_mask)
-
-    def decompose(self):
-        return self.tensors, self.mask
-
-    def __repr__(self):
-        return str(self.tensors)
-
-def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
-    # TODO make this more general
-    if tensor_list[0].ndim == 3:
-        if torchvision._is_tracing():
-            # nested_tensor_from_tensor_list() does not export well to ONNX
-            # call _onnx_nested_tensor_from_tensor_list() instead
-            return _onnx_nested_tensor_from_tensor_list(tensor_list)
-
-        # TODO make it support different-sized images
-        max_size = _max_by_axis([list(img.shape) for img in tensor_list])
-        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
-        batch_shape = [len(tensor_list)] + max_size
-        b, c, h, w = batch_shape
-        dtype = tensor_list[0].dtype
-        device = tensor_list[0].device
-        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
-        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
-        for img, pad_img, m in zip(tensor_list, tensor, mask):
-            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
-            m[: img.shape[1], : img.shape[2]] = False
-    elif tensor_list[0].ndim == 2:
-        if torchvision._is_tracing():
-            # nested_tensor_from_tensor_list() does not export well to ONNX
-            # call _onnx_nested_tensor_from_tensor_list() instead
-            return _onnx_nested_tensor_from_tensor_list(tensor_list)
-
-        # TODO make it support different-sized images
-        max_size = _max_by_axis([list(txt.shape) for txt in tensor_list])
-        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
-        batch_shape = [len(tensor_list)] + max_size
-        b, c, l = batch_shape
-        dtype = tensor_list[0].dtype
-        device = tensor_list[0].device
-        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
-        mask = torch.ones((b, l), dtype=torch.bool, device=device)
-        for txt, pad_txt, m in zip(tensor_list, tensor, mask):
-            pad_txt[: txt.shape[0], : txt.shape[1]] = txt
-            m[: txt.shape[1]] = False
-    else:
-        raise ValueError("not supported")
-    return NestedTensor(tensor, mask)
-
-def _collate_and_pad_divisibility(tensor_list: list, div=32):
-    max_size = []
-    for i in range(tensor_list[0].dim()):
-        max_size_i = torch.max(
-            torch.tensor([img.shape[i] for img in tensor_list]).to(torch.float32)
-        ).to(torch.int64)
-        max_size.append(max_size_i)
-    max_size = tuple(max_size)
-
-    c,h,w = max_size
-    pad_h = (div - h % div) if h % div != 0 else 0
-    pad_w = (div - w % div) if w % div != 0 else 0
-    max_size = (c,h+pad_h,w+pad_w)
-    
-    # work around for
-    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
-    # m[: img.shape[1], :img.shape[2]] = False
-    # which is not yet supported in onnx
-    padded_imgs = []
-    padded_masks = []
-    for img in tensor_list:
-        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
-        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
-        padded_imgs.append(padded_img)
-
-        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)
-        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)
-        padded_masks.append(padded_mask.to(torch.bool))
-    
-    return padded_imgs
-
-# _onnx_nested_tensor_from_tensor_list() is an implementation of
-# nested_tensor_from_tensor_list() that is supported by ONNX tracing.
-@torch.jit.unused
-def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:
-    max_size = []
-    for i in range(tensor_list[0].dim()):
-        max_size_i = torch.max(
-            torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)
-        ).to(torch.int64)
-        max_size.append(max_size_i)
-    max_size = tuple(max_size)
-
-    # work around for
-    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
-    # m[: img.shape[1], :img.shape[2]] = False
-    # which is not yet supported in onnx
-    padded_imgs = []
-    padded_masks = []
-    for img in tensor_list:
-        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
-        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
-        padded_imgs.append(padded_img)
-
-        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)
-        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)
-        padded_masks.append(padded_mask.to(torch.bool))
-
-    tensor = torch.stack(padded_imgs)
-    mask = torch.stack(padded_masks)
-
-    return NestedTensor(tensor, mask=mask)
-
-
-def is_dist_avail_and_initialized():
-    if not dist.is_available():
-        return False
-    if not dist.is_initialized():
-        return False
-    return True