5

analyze.py

@@ -0,0 +1,11 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import cv2
+import os
+
+with open('debug/stats.txt', 'r') as f:
+    stats = f.readlines()
+for stat in stats:
+    plt.scatter(float(stat.split(" ")[0]), float(stat.split(' ')[1]), alpha=0.1, s=10, c='red')
+    print(stat)
+plt.savefig('stats.png')

common/config.py

@@ -0,0 +1,31 @@
+r"""config"""
+import argparse
+
+def parse_opts():
+    r"""arguments"""
+    parser = argparse.ArgumentParser(description='Dense Cross-Query-and-Support Attention Weighted Mask Aggregation for Few-Shot Segmentation')
+
+    # common
+    parser.add_argument('--datapath', type=str, default='./datasets')
+    parser.add_argument('--benchmark', type=str, default='pascal', choices=['pascal', 'coco', 'fss'])
+    parser.add_argument('--fold', type=int, default=0, choices=[0, 1, 2, 3])
+    parser.add_argument('--bsz', type=int, default=20)
+    parser.add_argument('--nworker', type=int, default=8)
+    parser.add_argument('--backbone', type=str, default='swin', choices=['resnet50', 'resnet101', 'swin'])
+    parser.add_argument('--feature_extractor_path', type=str, default='')
+    parser.add_argument('--logpath', type=str, default='./logs')
+
+    # for train
+    parser.add_argument('--lr', type=float, default=1e-3)
+    parser.add_argument('--nepoch', type=int, default=1000)
+    parser.add_argument('--local-rank', default=0, type=int, help='node rank for distributed training')
+
+    # for test
+    parser.add_argument('--load', type=str, default='')
+    parser.add_argument('--nshot', type=int, default=1)
+    parser.add_argument('--visualize', action='store_true')
+    parser.add_argument('--vispath', type=str, default='./vis')
+    parser.add_argument('--use_original_imgsize', action='store_true')
+
+    args = parser.parse_args()
+    return args

common/evaluation.py

@@ -0,0 +1,39 @@
+r""" Evaluate mask prediction """
+import torch
+
+
+class Evaluator:
+    r""" Computes intersection and union between prediction and ground-truth """
+    @classmethod
+    def initialize(cls):
+        cls.ignore_index = 255
+
+    @classmethod
+    def classify_prediction(cls, pred_mask, batch):
+        gt_mask = batch.get('query_mask')
+
+        # Apply ignore_index in PASCAL-5i masks (following evaluation scheme in PFE-Net (TPAMI 2020))
+        query_ignore_idx = batch.get('query_ignore_idx')
+        if query_ignore_idx is not None:
+            assert torch.logical_and(query_ignore_idx, gt_mask).sum() == 0
+            query_ignore_idx *= cls.ignore_index
+            gt_mask = gt_mask + query_ignore_idx
+            pred_mask[gt_mask == cls.ignore_index] = cls.ignore_index
+
+        # compute intersection and union of each episode in a batch
+        area_inter, area_pred, area_gt = [],  [], []
+        for _pred_mask, _gt_mask in zip(pred_mask, gt_mask):
+            _inter = _pred_mask[_pred_mask == _gt_mask]
+            if _inter.size(0) == 0:  # as torch.histc returns error if it gets empty tensor (pytorch 1.5.1)
+                _area_inter = torch.tensor([0, 0], device=_pred_mask.device)
+            else:
+                _area_inter = torch.histc(_inter, bins=2, min=0, max=1)
+            area_inter.append(_area_inter)
+            area_pred.append(torch.histc(_pred_mask, bins=2, min=0, max=1))
+            area_gt.append(torch.histc(_gt_mask, bins=2, min=0, max=1))
+        area_inter = torch.stack(area_inter).t()
+        area_pred = torch.stack(area_pred).t()
+        area_gt = torch.stack(area_gt).t()
+        area_union = area_pred + area_gt - area_inter
+
+        return area_inter, area_union

common/logger.py

@@ -0,0 +1,117 @@
+r""" Logging during training/testing """
+import datetime
+import logging
+import os
+
+from tensorboardX import SummaryWriter
+import torch
+
+
+class AverageMeter:
+    r""" Stores loss, evaluation results """
+    def __init__(self, dataset):
+        self.benchmark = dataset.benchmark
+        self.class_ids_interest = dataset.class_ids
+        self.class_ids_interest = torch.tensor(self.class_ids_interest).cuda()
+
+        if self.benchmark == 'pascal':
+            self.nclass = 20
+        elif self.benchmark == 'coco':
+            self.nclass = 80
+        elif self.benchmark == 'fss':
+            self.nclass = 1000
+
+        self.intersection_buf = torch.zeros([2, self.nclass]).float().cuda()
+        self.union_buf = torch.zeros([2, self.nclass]).float().cuda()
+        self.ones = torch.ones_like(self.union_buf)
+        self.loss_buf = []
+
+    def update(self, inter_b, union_b, class_id, loss):
+        self.intersection_buf.index_add_(1, class_id, inter_b.float())
+        self.union_buf.index_add_(1, class_id, union_b.float())
+        if loss is None:
+            loss = torch.tensor(0.0)
+        self.loss_buf.append(loss)
+
+    def compute_iou(self):
+        iou = self.intersection_buf.float() / \
+              torch.max(torch.stack([self.union_buf, self.ones]), dim=0)[0]
+        iou = iou.index_select(1, self.class_ids_interest)
+        miou = iou[1].mean() * 100
+
+        fb_iou = (self.intersection_buf.index_select(1, self.class_ids_interest).sum(dim=1) /
+                  self.union_buf.index_select(1, self.class_ids_interest).sum(dim=1)).mean() * 100
+
+        return miou, fb_iou
+
+    def write_result(self, split, epoch):
+        iou, fb_iou = self.compute_iou()
+
+        loss_buf = torch.stack(self.loss_buf)
+        msg = '\n*** %s ' % split
+        msg += '[@Epoch %02d] ' % epoch
+        msg += 'Avg L: %6.5f  ' % loss_buf.mean()
+        msg += 'mIoU: %5.2f   ' % iou
+        msg += 'FB-IoU: %5.2f   ' % fb_iou
+
+        msg += '***\n'
+        Logger.info(msg)
+
+    def write_process(self, batch_idx, datalen, epoch, write_batch_idx=20):
+        if batch_idx % write_batch_idx == 0:
+            msg = '[Epoch: %02d] ' % epoch if epoch != -1 else ''
+            msg += '[Batch: %04d/%04d] ' % (batch_idx+1, datalen)
+            iou, fb_iou = self.compute_iou()
+            if epoch != -1:
+                loss_buf = torch.stack(self.loss_buf)
+                msg += 'L: %6.5f  ' % loss_buf[-1]
+                msg += 'Avg L: %6.5f  ' % loss_buf.mean()
+            msg += 'mIoU: %5.2f  |  ' % iou
+            msg += 'FB-IoU: %5.2f' % fb_iou
+            Logger.info(msg)
+
+
+class Logger:
+    r""" Writes evaluation results of training/testing """
+    @classmethod
+    def initialize(cls, args, training):
+        logtime = datetime.datetime.now().__format__('_%m%d_%H%M%S')
+        logpath = os.path.join(args.logpath, 'train/fold_' + str(args.fold) + logtime) if training \
+            else os.path.join(args.logpath, 'test/fold_' + args.load.split('/')[-2].split('.')[0] + logtime)
+        if logpath == '': logpath = logtime
+
+        cls.logpath = logpath
+        cls.benchmark = args.benchmark
+        if not os.path.exists(cls.logpath): os.makedirs(cls.logpath)
+
+        logging.basicConfig(filemode='w',
+                            filename=os.path.join(cls.logpath, 'log.txt'),
+                            level=logging.INFO,
+                            format='%(message)s',
+                            datefmt='%m-%d %H:%M:%S')
+
+        # Console log config
+        console = logging.StreamHandler()
+        console.setLevel(logging.INFO)
+        formatter = logging.Formatter('%(message)s')
+        console.setFormatter(formatter)
+        logging.getLogger('').addHandler(console)
+
+        # Tensorboard writer
+        cls.tbd_writer = SummaryWriter(os.path.join(cls.logpath, 'tbd/runs'))
+
+        # Log arguments
+        logging.info('\n:==================== Start =====================')
+        for arg_key in args.__dict__:
+            logging.info('| %20s: %-24s' % (arg_key, str(args.__dict__[arg_key])))
+        logging.info(':================================================\n')
+
+    @classmethod
+    def info(cls, msg):
+        r""" Writes log message to log.txt """
+        logging.info(msg)
+
+    @classmethod
+    def save_model_miou(cls, model, epoch, val_miou):
+        torch.save(model.state_dict(), os.path.join(cls.logpath, "model_{}.pt".format(epoch)))
+        cls.info('Model saved @%d w/ val. mIoU: %5.2f.\n' % (epoch, val_miou))

common/utils.py

@@ -0,0 +1,32 @@
+r""" Helper functions """
+import random
+
+import torch
+import numpy as np
+
+
+def fix_randseed(seed):
+    r""" Set random seeds for reproducibility """
+    if seed is None:
+        seed = int(random.random() * 1e5)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+
+def mean(x):
+    return sum(x) / len(x) if len(x) > 0 else 0.0
+
+
+def to_cuda(batch):
+    for key, value in batch.items():
+        if isinstance(value, torch.Tensor):
+            batch[key] = value.cuda()
+    return batch
+
+
+def to_cpu(tensor):
+    return tensor.detach().clone().cpu()

common/vis.py

@@ -0,0 +1,106 @@
+r""" Visualize model predictions """
+import os
+
+from PIL import Image
+import numpy as np
+import torchvision.transforms as transforms
+
+from . import utils
+
+
+class Visualizer:
+
+    @classmethod
+    def initialize(cls, visualize, vispath='./vis/'):
+        cls.visualize = visualize
+        if not visualize:
+            return
+
+        cls.colors = {'red': (255, 50, 50), 'blue': (102, 140, 255)}
+        for key, value in cls.colors.items():
+            cls.colors[key] = tuple([c / 255 for c in cls.colors[key]])
+
+        cls.mean_img = [0.485, 0.456, 0.406]
+        cls.std_img = [0.229, 0.224, 0.225]
+        cls.to_pil = transforms.ToPILImage()
+        cls.vis_path = vispath
+        if not os.path.exists(cls.vis_path): os.makedirs(cls.vis_path)
+
+    @classmethod
+    def visualize_prediction_batch(cls, spt_img_b, spt_mask_b, qry_img_b, qry_mask_b, pred_mask_b, cls_id_b, batch_idx, iou_b=None):
+        spt_img_b = utils.to_cpu(spt_img_b)
+        spt_mask_b = utils.to_cpu(spt_mask_b)
+        qry_img_b = utils.to_cpu(qry_img_b)
+        qry_mask_b = utils.to_cpu(qry_mask_b)
+        pred_mask_b = utils.to_cpu(pred_mask_b)
+        cls_id_b = utils.to_cpu(cls_id_b)
+
+        for sample_idx, (spt_img, spt_mask, qry_img, qry_mask, pred_mask, cls_id) in \
+                enumerate(zip(spt_img_b, spt_mask_b, qry_img_b, qry_mask_b, pred_mask_b, cls_id_b)):
+            iou = iou_b[sample_idx] if iou_b is not None else None
+            cls.visualize_prediction(spt_img, spt_mask, qry_img, qry_mask, pred_mask, cls_id, batch_idx, sample_idx, True, iou)
+
+    @classmethod
+    def to_numpy(cls, tensor, type):
+        if type == 'img':
+            return np.array(cls.to_pil(cls.unnormalize(tensor))).astype(np.uint8)
+        elif type == 'mask':
+            return np.array(tensor).astype(np.uint8)
+        else:
+            raise Exception('Undefined tensor type: %s' % type)
+
+    @classmethod
+    def visualize_prediction(cls, spt_imgs, spt_masks, qry_img, qry_mask, pred_mask, cls_id, batch_idx, sample_idx, label, iou=None):
+
+        spt_color = cls.colors['blue']
+        qry_color = cls.colors['red']
+        pred_color = cls.colors['red']
+
+        spt_imgs = [cls.to_numpy(spt_img, 'img') for spt_img in spt_imgs]
+        spt_pils = [cls.to_pil(spt_img) for spt_img in spt_imgs]
+        spt_masks = [cls.to_numpy(spt_mask, 'mask') for spt_mask in spt_masks]
+        spt_masked_pils = [Image.fromarray(cls.apply_mask(spt_img, spt_mask, spt_color)) for spt_img, spt_mask in zip(spt_imgs, spt_masks)]
+
+        qry_img = cls.to_numpy(qry_img, 'img')
+        qry_pil = cls.to_pil(qry_img)
+        qry_mask = cls.to_numpy(qry_mask, 'mask')
+        pred_mask = cls.to_numpy(pred_mask, 'mask')
+        pred_masked_pil = Image.fromarray(cls.apply_mask(qry_img.astype(np.uint8), pred_mask.astype(np.uint8), pred_color))
+        qry_masked_pil = Image.fromarray(cls.apply_mask(qry_img.astype(np.uint8), qry_mask.astype(np.uint8), qry_color))
+
+        merged_pil = cls.merge_image_pair(spt_masked_pils + [pred_masked_pil, qry_masked_pil])
+
+        iou = iou.item() if iou else 0.0
+        merged_pil.save(cls.vis_path + '%d_%d_class-%d_iou-%.2f' % (batch_idx, sample_idx, cls_id, iou) + '.jpg')
+
+    @classmethod
+    def merge_image_pair(cls, pil_imgs):
+        r""" Horizontally aligns a pair of pytorch tensor images (3, H, W) and returns PIL object """
+
+        canvas_width = sum([pil.size[0] for pil in pil_imgs])
+        canvas_height = max([pil.size[1] for pil in pil_imgs])
+        canvas = Image.new('RGB', (canvas_width, canvas_height))
+
+        xpos = 0
+        for pil in pil_imgs:
+            canvas.paste(pil, (xpos, 0))
+            xpos += pil.size[0]
+
+        return canvas
+
+    @classmethod
+    def apply_mask(cls, image, mask, color, alpha=0.5):
+        r""" Apply mask to the given image. """
+        for c in range(3):
+            image[:, :, c] = np.where(mask == 1,
+                                      image[:, :, c] *
+                                      (1 - alpha) + alpha * color[c] * 255,
+                                      image[:, :, c])
+        return image
+
+    @classmethod
+    def unnormalize(cls, img):
+        img = img.clone()
+        for im_channel, mean, std in zip(img, cls.mean_img, cls.std_img):
+            im_channel.mul_(std).add_(mean)
+        return img

gpu_mem_track.py

@@ -0,0 +1,113 @@
+import gc
+import datetime
+import inspect
+
+import torch
+import numpy as np
+
+dtype_memory_size_dict = {
+    torch.float64: 64/8,
+    torch.double: 64/8,
+    torch.float32: 32/8,
+    torch.float: 32/8,
+    torch.float16: 16/8,
+    torch.half: 16/8,
+    torch.int64: 64/8,
+    torch.long: 64/8,
+    torch.int32: 32/8,
+    torch.int: 32/8,
+    torch.int16: 16/8,
+    torch.short: 16/6,
+    torch.uint8: 8/8,
+    torch.int8: 8/8,
+}
+# compatibility of torch1.0
+if getattr(torch, "bfloat16", None) is not None:
+    dtype_memory_size_dict[torch.bfloat16] = 16/8
+if getattr(torch, "bool", None) is not None:
+    dtype_memory_size_dict[torch.bool] = 8/8 # pytorch use 1 byte for a bool, see https://github.com/pytorch/pytorch/issues/41571
+
+def get_mem_space(x):
+    try:
+        ret = dtype_memory_size_dict[x]
+    except KeyError:
+        print(f"dtype {x} is not supported!")
+    return ret
+
+class MemTracker(object):
+    """
+    Class used to track pytorch memory usage
+    Arguments:
+        detail(bool, default True): whether the function shows the detail gpu memory usage
+        path(str): where to save log file
+        verbose(bool, default False): whether show the trivial exception
+        device(int): GPU number, default is 0
+    """
+    def __init__(self, detail=True, path='', verbose=False, device=0):
+        self.print_detail = detail
+        self.last_tensor_sizes = set()
+        self.gpu_profile_fn = path + f'{datetime.datetime.now():%d-%b-%y-%H:%M:%S}-gpu_mem_track.txt'
+        self.verbose = verbose
+        self.begin = True
+        self.device = device
+
+    def get_tensors(self):
+        for obj in gc.get_objects():
+            try:
+                if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
+                    tensor = obj
+                else:
+                    continue
+                if tensor.is_cuda:
+                    yield tensor
+            except Exception as e:
+                if self.verbose:
+                    print('A trivial exception occured: {}'.format(e))
+
+    def get_tensor_usage(self):
+        sizes = [np.prod(np.array(tensor.size())) * get_mem_space(tensor.dtype) for tensor in self.get_tensors()]
+        return np.sum(sizes) / 1024**2
+
+    def get_allocate_usage(self):
+        return torch.cuda.memory_allocated() / 1024**2
+
+    def clear_cache(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def print_all_gpu_tensor(self, file=None):
+        for x in self.get_tensors():
+            print(x.size(), x.dtype, np.prod(np.array(x.size()))*get_mem_space(x.dtype)/1024**2, file=file)
+
+    def track(self):
+        """
+        Track the GPU memory usage
+        """
+        frameinfo = inspect.stack()[1]
+        where_str = frameinfo.filename + ' line ' + str(frameinfo.lineno) + ': ' + frameinfo.function
+
+        with open(self.gpu_profile_fn, 'a+') as f:
+
+            if self.begin:
+                f.write(f"GPU Memory Track | {datetime.datetime.now():%d-%b-%y-%H:%M:%S} |"
+                        f" Total Tensor Used Memory:{self.get_tensor_usage():<7.1f}Mb"
+                        f" Total Allocated Memory:{self.get_allocate_usage():<7.1f}Mb\n\n")
+                self.begin = False
+
+            if self.print_detail is True:
+                ts_list = [(tensor.size(), tensor.dtype) for tensor in self.get_tensors()]
+                new_tensor_sizes = {(type(x),
+                                    tuple(x.size()),
+                                    ts_list.count((x.size(), x.dtype)),
+                                    np.prod(np.array(x.size()))*get_mem_space(x.dtype)/1024**2,
+                                    x.dtype) for x in self.get_tensors()}
+                for t, s, n, m, data_type in new_tensor_sizes - self.last_tensor_sizes:
+                    f.write(f'+ | {str(n)} * Size:{str(s):<20} | Memory: {str(m*n)[:6]} M | {str(t):<20} | {data_type}\n')
+                for t, s, n, m, data_type in self.last_tensor_sizes - new_tensor_sizes:
+                    f.write(f'- | {str(n)} * Size:{str(s):<20} | Memory: {str(m*n)[:6]} M | {str(t):<20} | {data_type}\n')
+
+                self.last_tensor_sizes = new_tensor_sizes
+
+            f.write(f"\nAt {where_str:<50}"
+                    f" Total Tensor Used Memory:{self.get_tensor_usage():<7.1f}Mb"
+                    f" Total Allocated Memory:{self.get_allocate_usage():<7.1f}Mb\n\n")

importance_analysis.py

@@ -0,0 +1,130 @@
+r""" Dense Cross-Query-and-Support Attention Weighted Mask Aggregation for Few-Shot Segmentation """
+import torch.nn as nn
+import torch
+
+from model.DCAMA import DCAMA
+from common.logger import Logger, AverageMeter
+from common.vis import Visualizer
+from common.evaluation import Evaluator
+from common.config import parse_opts
+from common import utils
+from data.dataset import FSSDataset
+import cv2
+import numpy as np
+import os
+
+def test(model, dataloader, nshot):
+    r""" Test """
+
+    # Freeze randomness during testing for reproducibility
+    utils.fix_randseed(0)
+    average_meter = AverageMeter(dataloader.dataset)
+
+    for idx, batch in enumerate(dataloader):
+
+        # 1. forward pass
+        nshot = batch['support_imgs'].size(1)
+       	
+        batch['support_imgs'][0][0] = batch['query_img'][0]
+        batch['support_masks'][0][0] = batch['query_mask'][0]
+
+        batch = utils.to_cuda(batch)
+        pred_mask, simi, simi_map = model.module.predict_mask_nshot(batch, nshot=nshot)
+
+        assert pred_mask.size() == batch['query_mask'].size()
+
+        # 2. Evaluate prediction
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask.clone(), batch)
+        
+        ## TODO:
+        iou = area_inter[1] / area_union[1]
+        
+        '''
+        cv2.imwrite('debug/query.png', cv2.imread("/home/bkdongxianchi/MY_MOT/TWL/data/COCO2014/{}".format(batch['query_name'][0])))
+        cv2.imwrite('debug/query_mask.png', (batch['query_mask'][0] * 255).detach().cpu().numpy().astype(np.uint8))
+        cv2.imwrite('debug/support_{:.3}.png'.format(iou.item()), cv2.imread('/home/bkdongxianchi/MY_MOT/TWL/data/COCO2014/{}'.format(batch['support_names'][0][0])))
+        cv2.imwrite('debug/support_mask_{:.3}.png'.format(iou.item()), (batch['support_masks'][0][0] * 255).detach().cpu().numpy().astype(np.uint8))
+        simi_map = simi_map - simi_map.min()
+        simi_map = (simi_map / simi_map.max() * 255).detach().cpu().numpy().astype(np.uint8)
+        cv2.imwrite('debug/simi_map_{:.3}.png'.format(iou.item()), simi_map)
+        
+        if os.path.exists('debug/stats.txt'):
+            with open('debug/stats.txt', "a") as f:
+                f.write("{} {}\n".format(simi.item(), iou.item()))
+        else:
+            with open('debug/stats.txt', 'w') as f:
+                f.write('{} {}\n'.format(simi.item(), iou.item()))
+        '''
+
+        average_meter.update(area_inter, area_union, batch['class_id'], loss=None)
+        average_meter.write_process(idx, len(dataloader), epoch=-1, write_batch_idx=1)
+
+        # Visualize predictions
+        if Visualizer.visualize:
+            Visualizer.visualize_prediction_batch(batch['support_imgs'], batch['support_masks'],
+                                                  batch['query_img'], batch['query_mask'],
+                                                  pred_mask, batch['class_id'], idx,
+                                                  iou_b=area_inter[1].float() / area_union[1].float())
+
+    # Write evaluation results
+    average_meter.write_result('Test', 0)
+    miou, fb_iou = average_meter.compute_iou()
+
+    return miou, fb_iou
+
+
+if __name__ == '__main__':
+
+    # Arguments parsing
+    args = parse_opts()
+
+    Logger.initialize(args, training=False)
+
+    # Model initialization
+    model = DCAMA(args.backbone, args.feature_extractor_path, args.use_original_imgsize)
+    model.eval()
+
+    # Device setup
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    Logger.info('# available GPUs: %d' % torch.cuda.device_count())
+    model = nn.DataParallel(model)
+    model.to(device)
+
+    # Load trained model
+    if args.load == '': raise Exception('Pretrained model not specified.')
+    params = model.state_dict()
+    state_dict = torch.load(args.load)
+    
+    if 'state_dict' in state_dict.keys():
+        state_dict = state_dict['state_dict']
+
+    state_dict2 = {}
+    for k, v in state_dict.items():
+        if 'scorer' not in k:
+            state_dict2[k] = v
+    state_dict = state_dict2
+    for k1, k2 in zip(list(state_dict.keys()), params.keys()):
+        state_dict[k2] = state_dict.pop(k1)
+    
+
+    try:
+        model.load_state_dict(state_dict, strict=True)
+    except:
+        for k in params.keys():
+            if k not in state_dict.keys():
+                state_dict[k] = params[k]
+        model.load_state_dict(state_dict, strict=True)
+
+    # Helper classes (for testing) initialization
+    Evaluator.initialize()
+    Visualizer.initialize(args.visualize, args.vispath)
+
+    # Dataset initialization
+    FSSDataset.initialize(img_size=384, datapath=args.datapath, use_original_imgsize=args.use_original_imgsize)
+    dataloader_test = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'test', args.nshot)
+
+    # Test
+    with torch.no_grad():
+        test_miou, test_fb_iou = test(model, dataloader_test, args.nshot)
+    Logger.info('Fold %d mIoU: %5.2f \t FB-IoU: %5.2f' % (args.fold, test_miou.item(), test_fb_iou.item()))
+    Logger.info('==================== Finished Testing ====================')

model/DCAMA.py

@@ -0,0 +1,625 @@
+r""" Dense Cross-Query-and-Support Attention Weighted Mask Aggregation for Few-Shot Segmentation """
+from functools import reduce
+from operator import add
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import resnet
+
+from .base.swin_transformer import SwinTransformer
+from model.base.transformer import MultiHeadedAttention, PositionalEncoding
+import copy
+
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), x.size(1), -1).contiguous()
+
+def reshape(x, size):
+    size1 = torch.tensor(x.size()).float().cuda()
+    # x = torch.logical_not(x.cuda())
+    yxs = torch.stack(torch.where(x), dim=-1)
+    ratio = size[0] / size1[0]
+    yxs2 = (yxs * ratio).long()
+    x2 = torch.zeros((size[0], size[1])).float().cuda()
+    return yxs2
+
+
+class DCAMA(nn.Module):
+
+    def __init__(self, backbone, pretrained_path, use_original_imgsize):
+        super(DCAMA, self).__init__()
+
+        self.backbone = backbone
+        self.use_original_imgsize = use_original_imgsize
+
+        # feature extractor initialization
+        if backbone == 'resnet50':
+            self.feature_extractor = resnet.resnet50()
+            self.feature_extractor.load_state_dict(torch.load(pretrained_path))
+            self.feat_channels = [256, 512, 1024, 2048]
+            self.nlayers = [3, 4, 6, 3]
+            self.feat_ids = list(range(0, 17))
+            self.last_feat_size = [12, 12]
+        elif backbone == 'resnet101':
+            self.feature_extractor = resnet.resnet101()
+            self.feature_extractor.load_state_dict(torch.load(pretrained_path))
+            self.feat_channels = [256, 512, 1024, 2048]
+            self.nlayers = [3, 4, 23, 3]
+            self.feat_ids = list(range(0, 34))
+        elif backbone == 'swin':
+            self.feature_extractor = SwinTransformer(img_size=384, patch_size=4, window_size=12, embed_dim=128,
+                                            depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32])
+            self.feature_extractor.load_state_dict(torch.load(pretrained_path)['model'])
+            self.feat_channels = [128, 256, 512, 1024]
+            self.nlayers = [2, 2, 18, 2]
+        else:
+            raise Exception('Unavailable backbone: %s' % backbone)
+        self.feature_extractor.eval()
+
+        # define model
+        self.lids = reduce(add, [[i + 1] * x for i, x in enumerate(self.nlayers)])
+        self.stack_ids = torch.tensor(self.lids).bincount()[-4:].cumsum(dim=0)
+        self.model = DCAMA_model(in_channels=self.feat_channels, stack_ids=self.stack_ids)
+        
+        ## TODO:
+        
+        self.scorer2 = nn.ModuleList()
+        for layer_idx in range(len(self.nlayers)):
+            layer_num = self.nlayers[layer_idx]
+            for idx in range(layer_num):
+                self.scorer2.append(
+                    nn.Sequential(
+                        nn.Conv2d(256 * 2 ** layer_idx, 256 * 2 ** layer_idx, 1, 1),
+                        # nn.ReLU(),
+                        # nn.InstanceNorm2d(256 * 2 ** layer_idx),
+                        # nn.Conv2d(256 * 2 ** layer_idx, 256 * 2 ** layer_idx, 1, 1),
+                    )
+                )
+        self.scorer1 = nn.Sequential(
+            nn.Linear(sum(self.nlayers) - self.nlayers[0], 1)    
+        )
+       
+        self.cross_entropy_loss = nn.CrossEntropyLoss()
+
+    def forward(self, query_img, support_img, support_mask, nshot, predict_score=False):
+        n_support_feats = []
+        with torch.no_grad():
+            for k in range(nshot):
+                support_feats_= self.extract_feats(support_img[:, k])
+                support_feats = copy.deepcopy(support_feats_)
+                del support_feats_
+                torch.cuda.empty_cache()
+                n_support_feats.append(support_feats)
+            query_feats = self.extract_feats(query_img)
+        
+
+        logit_mask = self.model(query_feats, n_support_feats, support_mask.clone(), nshot=nshot)
+        ## TODO:
+        MAX_SHOTS = 1
+        if len(n_support_feats) >= MAX_SHOTS:
+            nshot = MAX_SHOTS   
+            n_support_query_f = []
+            n_simi = []
+            for i in range(len(n_support_feats)):
+                support_f = n_support_feats[i]
+                support_query_f = []
+                simi_l = []
+                simi_sum = []
+                for l in range(len(query_feats)):
+                    if l < self.stack_ids[0]:
+                        continue
+                    elif l < self.stack_ids[1]:
+                        DCAMA_blocks = self.model.DCAMA_blocks[0]
+                        pe = self.model.pe[0]
+                    elif l < self.stack_ids[2]:
+                        DCAMA_blocks = self.model.DCAMA_blocks[1]
+                        pe = self.model.pe[1]
+                    else:
+                        DCAMA_blocks = self.model.DCAMA_blocks[2]
+                        pe = self.model.pe[2]
+                    a_support_f = support_f[l].clone()
+                    coords = reshape(support_mask[0, i], a_support_f.size()[-2:])
+                    b, ch, w, h = a_support_f.size()
+                    a_support_f = a_support_f.view(b, ch, -1)
+                    a_support_f = DCAMA_blocks.linears[0](pe(a_support_f.permute(0, 2, 1))).permute(0, 2, 1)
+                    a_support_f = a_support_f.view(b, ch, w, h)
+                    a_support_f = self.scorer2[l](a_support_f)
+                    a_support_f = a_support_f[:, :, coords[:, 0], coords[:, 1]].mean(-1).unsqueeze(-1).unsqueeze(-1).repeat((1, 1, a_support_f.size(-2), a_support_f.size(-1)))
+                    # a_support_f[:, :, coords_reverse[:, 0], coords_reverse[:, 1]] *= 0.
+                    query_feat = query_feats[l].view(b, ch, -1)
+                    query_feat = DCAMA_blocks.linears[0](pe(query_feat.permute(0, 2, 1))).permute(0, 2, 1)
+                    query_feat = query_feat.view(b, ch, w, h)
+                    query_feat = self.scorer2[l](query_feat)
+                    simi = ((query_feat * a_support_f).sum(1)  / torch.norm(query_feat, dim=1) / torch.norm(a_support_f, dim=1))[0]
+                    simi_sum.append(simi)
+                    # simi = torch.norm(query_feats[l] - a_support_f, dim=1)[0]
+                    if l == 6:
+                        simi_map = simi.clone()
+                    simi = simi.view(-1).mean()
+                    simi_l.append(simi)
+                # simi_l = self.scorer1(torch.stack(simi_l, dim=0).unsqueeze(0)).squeeze(0)[0]
+                n_simi.append(torch.stack(simi_l, dim=0).mean())
+             
+            n_simi = torch.stack(n_simi, dim=0)
+            args = n_simi.argsort(descending=True)[:MAX_SHOTS]
+            support_mask = support_mask[:, args, :, :]
+            # n_support_feats = [n_support_feats[arg] for arg in args]
+            n_simis = n_simi[args].max()
+        else:
+            n_simis = torch.tensor(0.).float().cuda()
+        return logit_mask, n_simis
+    
+    def extract_feats(self, img):
+        r""" Extract input image features """
+        feats = []
+
+        if self.backbone == 'swin':
+            _ = self.feature_extractor.forward_features(img)
+            for feat in self.feature_extractor.feat_maps:
+                bsz, hw, c = feat.size()
+                h = int(hw ** 0.5)
+                feat = feat.view(bsz, h, h, c).permute(0, 3, 1, 2).contiguous()
+                feats.append(feat)
+        elif self.backbone == 'resnet50' or self.backbone == 'resnet101':
+            bottleneck_ids = reduce(add, list(map(lambda x: list(range(x)), self.nlayers)))
+            # Layer 0
+            feat = self.feature_extractor.conv1.forward(img)
+            feat = self.feature_extractor.bn1.forward(feat)
+            feat = self.feature_extractor.relu.forward(feat)
+            feat = self.feature_extractor.maxpool.forward(feat)
+
+            # Layer 1-4
+            for hid, (bid, lid) in enumerate(zip(bottleneck_ids, self.lids)):
+                res = feat
+                feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].conv1.forward(feat)
+                feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].bn1.forward(feat)
+                feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].relu.forward(feat)
+                feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].conv2.forward(feat)
+                feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].bn2.forward(feat)
+                feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].relu.forward(feat)
+                feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].conv3.forward(feat)
+                feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].bn3.forward(feat)
+
+                if bid == 0:
+                    res = self.feature_extractor.__getattr__('layer%d' % lid)[bid].downsample.forward(res)
+
+                feat += res
+
+                if hid + 1 in self.feat_ids:
+                    feats.append(feat.clone())
+
+                feat = self.feature_extractor.__getattr__('layer%d' % lid)[bid].relu.forward(feat)
+
+        return feats
+
+    def predict_mask_nshot(self, batch, nshot):
+        r""" n-shot inference """
+        query_img = batch['query_img']
+        support_imgs = batch['support_imgs']
+        support_masks = batch['support_masks']
+
+        if nshot == 1:
+            with torch.no_grad():
+                query_feats = self.extract_feats(query_img)
+                n_support_feats = []
+                for k in range(nshot):
+                    support_feats = self.extract_feats(support_imgs[:, k])
+                    n_support_feats.append(support_feats)
+
+            n_simis = []
+            simi_map = None
+            for i in range(len(n_support_feats)):
+                support_f = n_support_feats[i]
+                support_query_f = []
+                simi_l = []
+                for l in range(len(query_feats)):
+                    if l < self.stack_ids[0]:
+                        continue
+                    elif l < self.stack_ids[1]:
+                        DCAMA_blocks = self.model.DCAMA_blocks[0]
+                        pe = self.model.pe[0]
+                    elif l < self.stack_ids[2]:
+                        DCAMA_blocks = self.model.DCAMA_blocks[1]
+                        pe = self.model.pe[1]
+                    else:
+                        DCAMA_blocks = self.model.DCAMA_blocks[2]
+                        pe = self.model.pe[2]
+                    a_support_f = support_f[l].clone()
+                    coords = reshape(support_masks[0, i], a_support_f.size()[-2:])
+                    b, ch, w, h = a_support_f.size()
+                    a_support_f = a_support_f.view(b, ch, -1)
+                    a_support_f = DCAMA_blocks.linears[0](pe(a_support_f.permute(0, 2, 1))).permute(0, 2, 1)
+                    a_support_f = a_support_f.view(b, ch, w, h)
+                    a_support_f = a_support_f[:, :, coords[:, 0], coords[:, 1]].mean(-1).unsqueeze(-1).unsqueeze(-1).repeat((1, 1, a_support_f.size(-2), a_support_f.size(-1)))
+                    # a_support_f[:, :, coords_reverse[:, 0], coords_reverse[:, 1]] *= 0.
+                    query_feat = query_feats[l].view(b, ch, -1)
+                    query_feat = DCAMA_blocks.linears[0](pe(query_feat.permute(0, 2, 1))).permute(0, 2, 1)
+                    query_feat = query_feat.view(b, ch, w, h)
+                    simi = ((query_feat * a_support_f).sum(1)  / torch.norm(query_feat, dim=1) / torch.norm(a_support_f, dim=1))[0]
+                    # simi = torch.norm(query_feats[l] - a_support_f, dim=1)[0]
+                    if l == 13:
+                        simi_map = simi.clone()
+                    simi = simi.view(-1).max()
+                    simi_l.append(simi)
+                simi_l = torch.stack(simi_l, dim=0).mean()
+                n_simis.append(simi_l)
+            n_simis = torch.stack(n_simis, dim=0)
+            logit_mask = self.model(query_feats, n_support_feats, support_masks.clone(), nshot)
+            
+        else:
+            with torch.no_grad():
+                query_feats = self.extract_feats(query_img)
+                n_support_feats = []
+                for k in range(nshot):
+                    support_feats = self.extract_feats(support_imgs[:, k])
+                    n_support_feats.append(support_feats)
+            
+            ## TODO: retrieval V1 ##
+            MAX_SHOTS = 200
+            '''
+            if len(n_support_feats) > MAX_SHOTS:
+                nshot = MAX_SHOTS   
+                n_support_query_f = []
+                n_simis = []
+                for i in range(len(n_support_feats)):
+                    support_f = n_support_feats[i]
+                    support_query_f = []
+                    simi_l = []
+                    simi_sum = []
+                    for l in range(len(query_feats)):
+                        if l < self.stack_ids[0]:
+                            continue
+                        elif l < self.stack_ids[1]:
+                            DCAMA_blocks = self.model.DCAMA_blocks[0]
+                            pe = self.model.pe[0]
+                        elif l < self.stack_ids[2]:
+                            DCAMA_blocks = self.model.DCAMA_blocks[1]
+                            pe = self.model.pe[1]
+                        else:
+                            DCAMA_blocks = self.model.DCAMA_blocks[2]
+                            pe = self.model.pe[2]
+                        a_support_f = support_f[l].clone()
+                        coords = reshape(support_masks[0, i], a_support_f.size()[-2:])
+                        b, ch, w, h = a_support_f.size()
+                        a_support_f = a_support_f.view(b, ch, -1)
+                        a_support_f = DCAMA_blocks.linears[0](pe(a_support_f.permute(0, 2, 1))).permute(0, 2, 1)
+                        a_support_f = a_support_f.view(b, ch, w, h)
+                        a_support_f = a_support_f[:, :, coords[:, 0], coords[:, 1]].mean(-1).unsqueeze(-1).unsqueeze(-1).repeat((1, 1, a_support_f.size(-2), a_support_f.size(-1)))
+                        # a_support_f[:, :, coords_reverse[:, 0], coords_reverse[:, 1]] *= 0.
+                        query_feat = query_feats[l].view(b, ch, -1)
+                        query_feat = DCAMA_blocks.linears[0](pe(query_feat.permute(0, 2, 1))).permute(0, 2, 1)
+                        query_feat = query_feat.view(b, ch, w, h)
+                        simi = ((query_feat * a_support_f).sum(1)  / torch.norm(query_feat, dim=1) / torch.norm(a_support_f, dim=1))[0]
+                        simi_sum.append(simi)
+                        # simi = torch.norm(query_feats[l] - a_support_f, dim=1)[0]
+                        if l == 6:
+                            simi_map = simi.clone().detach().cpu().numpy()
+                        simi = simi.view(-1).max()
+                        simi_l.append(simi)
+                    simi_l = torch.stack(simi_l, dim=0).mean()
+                    n_simis.append(simi_l)
+                n_simis = torch.stack(n_simis, dim=0)
+                # nshot = max((n_simis > 0.).sum(), 1)
+                nshot = len(n_simis)
+                support_masks = support_masks[:, n_simis.argsort(descending=True)[:nshot], :, :]
+                n_support_feats = [n_support_feats[i] for i in n_simis.argsort(descending=True)[:nshot]]
+            else:
+                n_simis = torch.tensor(0.).float().cuda()
+                simi_map = None
+            ## TODO: retriever V2
+            '''
+            '''
+            MAX_SHOTS = 30
+            if len(n_support_feats) > MAX_SHOTS:
+                nshot = MAX_SHOTS
+                n_support_query_f = []
+                n_simis = []
+                support_f_list = []
+                n_support_feats2 = []
+                query_feats2 = []
+                for i in range(len(n_support_feats)):
+                    support_f = n_support_feats[i]
+                    n_support_feats2_l = []
+                    query_feats2_l = []
+                    for l in range(len(query_feats)):
+                        if l < self.stack_ids[0]:
+                            continue
+                        elif l < self.stack_ids[1]:
+                            DCAMA_blocks = self.model.DCAMA_blocks[0]
+                            pe = self.model.pe[0]
+                        elif l < self.stack_ids[2]:
+                            DCAMA_blocks = self.model.DCAMA_blocks[1]
+                            pe = self.model.pe[1]
+                        else:
+                            DCAMA_blocks = self.model.DCAMA_blocks[2]
+                            pe = self.model.pe[2]
+                        a_support_f = support_f[l].clone()
+                        coords = reshape(support_masks[0, i], a_support_f.size()[-2:])
+                        b, ch, w, h = a_support_f.size()
+                        a_support_f = a_support_f.view(b, ch, -1)
+                        a_support_f = DCAMA_blocks.linears[0](pe(a_support_f.permute(0, 2, 1))).permute(0, 2, 1)
+                        a_support_f = a_support_f.view(b, ch, w, h)
+                        a_support_f = a_support_f[:, :, coords[:, 0], coords[:, 1]].mean(-1)
+                        n_support_feats2_l.append(a_support_f)
+                        query_feat = query_feats[l].view(b, ch, -1)
+                        query_feat = DCAMA_blocks.linears[0](pe(query_feat.permute(0, 2, 1))).permute(0, 2, 1)
+                        query_feat = query_feat.view(b, ch, w, h)
+                        query_feats2_l.append(query_feat)
+                    n_support_feats2.append(n_support_feats2_l)
+                    query_feats2.append(query_feats2_l)
+                n_support_feats3 = [[] for _ in range(len(query_feats2[0]))]
+                selected = []
+                for i in range(MAX_SHOTS):
+                    simi_min = -100
+                    idx_min = -1
+                    for idx in range(len(n_support_feats2)):
+                        if idx in selected:
+                            continue
+                        support_feats2 = n_support_feats2[idx]
+                        simi = []
+                        for l in range(len(query_feats2[i])):
+                            support_feats_avg = torch.stack(n_support_feats3[l] + [support_feats2[l]], dim=0).mean(0)
+                            query_feat = query_feats2[i][l]
+                            a_support_f = support_feats_avg.unsqueeze(-1).unsqueeze(-1).repeat(
+                                (1, 1, query_feat.size(-2), query_feat.size(-1)))
+                            simi_l = ((query_feat * a_support_f).sum(1) / torch.norm(query_feat, dim=1) / torch.norm(
+                                a_support_f, dim=1))[0].view(-1).max()
+                            simi.append(simi_l)
+                        simi = torch.stack(simi, dim=0).mean()
+                        if simi > simi_min:
+                            simi_min = simi
+                            idx_min = idx
+                    support_feats2_argmin = n_support_feats2[idx]
+                    for l2 in range(len(query_feats2[0])):
+                        n_support_feats3[l2].append(n_support_feats2[idx_min][l2])
+                    selected.append(idx_min)
+                n_support_feats4 = []
+                for idx in selected:
+                    n_support_feats4.append(n_support_feats[idx])
+                support_masks = support_masks[:, torch.tensor(selected).long().cuda(), :, :]
+                n_support_feats = n_support_feats4
+                simi_map = None
+            else:
+                n_simis = torch.tensor(0.).float().cuda()
+                simi_map = None
+            '''
+            ## TODO: v3
+            
+            MAX_SHOTS = 200
+            if len(n_support_feats) > MAX_SHOTS:
+                nshot = MAX_SHOTS
+                n_support_query_f = []
+                n_simis = []
+                support_f_list = []
+                n_support_feats2 = []
+                query_feats2 = []
+                for i in range(len(n_support_feats)):
+                    support_f = n_support_feats[i]
+                    n_support_feats2_l = []
+                    query_feats2_l = []
+                    for l in range(len(query_feats)):
+                        if l < self.stack_ids[0]:
+                            continue
+                        elif l < self.stack_ids[1]:
+                            DCAMA_blocks = self.model.DCAMA_blocks[0]
+                            pe = self.model.pe[0]
+                        elif l < self.stack_ids[2]:
+                            DCAMA_blocks = self.model.DCAMA_blocks[1]
+                            pe = self.model.pe[1]
+                        else:
+                            DCAMA_blocks = self.model.DCAMA_blocks[2]
+                            pe = self.model.pe[2]
+                        a_support_f = support_f[l].clone()
+                        coords = reshape(support_masks[0, i], a_support_f.size()[-2:])
+                        b, ch, w, h = a_support_f.size()
+                        a_support_f = a_support_f.view(b, ch, -1)
+                        a_support_f_tmp = DCAMA_blocks.linears[0](pe(a_support_f.permute(0, 2, 1))).permute(0, 2, 1)
+                        a_support_f = a_support_f_tmp / a_support_f_tmp.norm(dim=1, keepdim=True) * DCAMA_blocks.linears[1](pe(a_support_f.permute(0, 2, 1))).permute(0, 2, 1)
+                        a_support_f = a_support_f.view(b, ch, w, h)
+                        a_support_f = a_support_f[:, :, coords[:, 0], coords[:, 1]].mean(-1)
+                        n_support_feats2_l.append(a_support_f)
+                        query_feat = query_feats[l].view(b, ch, -1)
+                        query_feat_tmp = DCAMA_blocks.linears[0](pe(query_feat.permute(0, 2, 1))).permute(0, 2, 1)
+                        query_feat = query_feat_tmp / query_feat_tmp.norm(dim=1, keepdim=True) * DCAMA_blocks.linears[1](pe(query_feat.permute(0, 2, 1))).permute(0, 2, 1)
+                        query_feat = query_feat.view(b, ch, w, h)
+                        query_feats2_l.append(query_feat)
+                    n_support_feats2.append(n_support_feats2_l)
+                    query_feats2.append(query_feats2_l)
+                n_support_feats3 = [[] for _ in range(len(query_feats2[0]))]
+                selected = []
+                for i in range(MAX_SHOTS):
+                    simi_min = -100
+                    idx_min = -1
+                    for idx in range(len(n_support_feats2)):
+                        if idx in selected:
+                            continue
+                        support_feats2 = n_support_feats2[idx]
+                        simi = []
+                        for l in range(len(query_feats2[i])):
+                            support_feats_avg = torch.stack(n_support_feats3[l] + [support_feats2[l]], dim=0).mean(0)
+                            query_feat = query_feats2[i][l]
+                            a_support_f = support_feats_avg.unsqueeze(-1).unsqueeze(-1).repeat(
+                                (1, 1, query_feat.size(-2), query_feat.size(-1)))
+                            simi_l = ((query_feat * a_support_f).sum(1))[0].view(-1).max()
+                            simi.append(simi_l)
+                        simi = torch.stack(simi, dim=0).mean()
+                        if simi > simi_min:
+                            simi_min = simi
+                            idx_min = idx
+                    support_feats2_argmin = n_support_feats2[idx]
+                    for l2 in range(len(query_feats2[0])):
+                        n_support_feats3[l2].append(n_support_feats2[idx_min][l2])
+                    selected.append(idx_min)
+                n_support_feats4 = []
+                for idx in selected:
+                    n_support_feats4.append(n_support_feats[idx])
+                support_masks = support_masks[:, torch.tensor(selected).long().cuda(), :, :]
+                n_support_feats = n_support_feats4
+                simi_map = None
+            else:
+                n_simis = torch.tensor(0.).float().cuda()
+                simi_map = None
+            
+            logit_mask = self.model(query_feats, n_support_feats, support_masks.clone(), nshot)
+
+        if self.use_original_imgsize:
+            org_qry_imsize = tuple([batch['org_query_imsize'][1].item(), batch['org_query_imsize'][0].item()])
+            logit_mask = F.interpolate(logit_mask, org_qry_imsize, mode='bilinear', align_corners=True)
+        else:
+            logit_mask = F.interpolate(logit_mask, support_imgs[0].size()[2:], mode='bilinear', align_corners=True)
+    
+        return logit_mask.argmax(dim=1), n_simis, simi_map
+
+    def compute_objective(self, logit_mask, gt_mask):
+        bsz = logit_mask.size(0)
+        logit_mask = logit_mask.view(bsz, 2, -1)
+        gt_mask = gt_mask.view(bsz, -1).long()
+
+        return self.cross_entropy_loss(logit_mask, gt_mask)
+
+    def train_mode(self):
+        self.train()
+        self.feature_extractor.eval()
+
+
+class DCAMA_model(nn.Module):
+    def __init__(self, in_channels, stack_ids):
+        super(DCAMA_model, self).__init__()
+
+        self.stack_ids = stack_ids
+
+        # DCAMA blocks
+        self.DCAMA_blocks = nn.ModuleList()
+        self.pe = nn.ModuleList()
+        for inch in in_channels[1:]:
+            self.DCAMA_blocks.append(MultiHeadedAttention(h=8, d_model=inch, dropout=0.5))
+            self.pe.append(PositionalEncoding(d_model=inch, dropout=0.5))
+
+        outch1, outch2, outch3 = 16, 64, 128
+
+        # conv blocks
+        self.conv1 = self.build_conv_block(stack_ids[3]-stack_ids[2], [outch1, outch2, outch3], [3, 3, 3], [1, 1, 1]) # 1/32
+        self.conv2 = self.build_conv_block(stack_ids[2]-stack_ids[1], [outch1, outch2, outch3], [5, 3, 3], [1, 1, 1]) # 1/16
+        self.conv3 = self.build_conv_block(stack_ids[1]-stack_ids[0], [outch1, outch2, outch3], [5, 5, 3], [1, 1, 1]) # 1/8
+
+        self.conv4 = self.build_conv_block(outch3, [outch3, outch3, outch3], [3, 3, 3], [1, 1, 1]) # 1/32 + 1/16
+        self.conv5 = self.build_conv_block(outch3, [outch3, outch3, outch3], [3, 3, 3], [1, 1, 1]) # 1/16 + 1/8
+
+        # mixer blocks
+        self.mixer1 = nn.Sequential(nn.Conv2d(outch3+2*in_channels[1]+2*in_channels[0], outch3, (3, 3), padding=(1, 1), bias=True),
+                                      nn.ReLU(),
+                                      nn.Conv2d(outch3, outch2, (3, 3), padding=(1, 1), bias=True),
+                                      nn.ReLU())
+
+        self.mixer2 = nn.Sequential(nn.Conv2d(outch2, outch2, (3, 3), padding=(1, 1), bias=True),
+                                      nn.ReLU(),
+                                      nn.Conv2d(outch2, outch1, (3, 3), padding=(1, 1), bias=True),
+                                      nn.ReLU())
+
+        self.mixer3 = nn.Sequential(nn.Conv2d(outch1, outch1, (3, 3), padding=(1, 1), bias=True),
+                                      nn.ReLU(),
+                                      nn.Conv2d(outch1, 2, (3, 3), padding=(1, 1), bias=True))
+    
+        
+
+    def forward(self, query_feats, support_feats, support_mask, nshot=1):
+        coarse_masks = []
+        for idx, query_feat in enumerate(query_feats):
+            # 1/4 scale feature only used in skip connect
+            if idx < self.stack_ids[0]: continue
+
+            bsz, ch, ha, wa = query_feat.size()
+
+            # reshape the input feature and mask
+            query = query_feat.view(bsz, ch, -1).permute(0, 2, 1).contiguous()
+            # if nshot == 1:
+            #     support_feat = support_feats[idx]
+            #     mask = F.interpolate(support_mask.unsqueeze(1).float(), support_feat.size()[2:], mode='bilinear',
+            #                          align_corners=True).view(support_feat.size()[0], -1)
+            #     support_feat = support_feat.view(support_feat.size()[0], support_feat.size()[1], -1).permute(0, 2, 1).contiguous()
+            # else:
+            support_feat = torch.stack([support_feats[k][idx] for k in range(nshot)])
+            support_feat = support_feat.view(-1, ch, ha * wa).permute(0, 2, 1).contiguous()
+            mask = torch.stack([F.interpolate(k.unsqueeze(1).float(), (ha, wa), mode='bilinear', align_corners=True)
+                                    for k in support_mask])
+            mask = mask.view(bsz, -1)
+
+            # DCAMA blocks forward
+            DCAMA_blocks = None
+            pe = None
+            if idx < self.stack_ids[1]:
+                DCAMA_blocks = self.DCAMA_blocks[0]
+                pe = self.pe[0]
+            elif idx < self.stack_ids[2]:
+                DCAMA_blocks = self.DCAMA_blocks[1]
+                pe = self.pe[1]
+            else:
+                DCAMA_blocks = self.DCAMA_blocks[2]
+                pe = self.pe[2]
+            coarse_mask = DCAMA_blocks(pe(query), pe(support_feat), mask)
+            coarse_masks.append(coarse_mask.permute(0, 2, 1).contiguous().view(bsz, 1, ha, wa))
+
+
+        # multi-scale conv blocks forward
+        bsz, ch, ha, wa = coarse_masks[self.stack_ids[3]-1-self.stack_ids[0]].size()
+        coarse_masks1 = torch.stack(coarse_masks[self.stack_ids[2]-self.stack_ids[0]:self.stack_ids[3]-self.stack_ids[0]]).transpose(0, 1).contiguous().view(bsz, -1, ha, wa)
+        bsz, ch, ha, wa = coarse_masks[self.stack_ids[2]-1-self.stack_ids[0]].size()
+        coarse_masks2 = torch.stack(coarse_masks[self.stack_ids[1]-self.stack_ids[0]:self.stack_ids[2]-self.stack_ids[0]]).transpose(0, 1).contiguous().view(bsz, -1, ha, wa)
+        bsz, ch, ha, wa = coarse_masks[self.stack_ids[1]-1-self.stack_ids[0]].size()
+        coarse_masks3 = torch.stack(coarse_masks[0:self.stack_ids[1]-self.stack_ids[0]]).transpose(0, 1).contiguous().view(bsz, -1, ha, wa)
+
+        coarse_masks1 = self.conv1(coarse_masks1)
+        coarse_masks2 = self.conv2(coarse_masks2)
+        coarse_masks3 = self.conv3(coarse_masks3)
+
+        # multi-scale cascade (pixel-wise addition)
+        coarse_masks1 = F.interpolate(coarse_masks1, coarse_masks2.size()[-2:], mode='bilinear', align_corners=True)
+        mix = coarse_masks1 + coarse_masks2
+        mix = self.conv4(mix)
+
+        mix = F.interpolate(mix, coarse_masks3.size()[-2:], mode='bilinear', align_corners=True)
+        mix = mix + coarse_masks3
+        mix = self.conv5(mix)
+
+        # skip connect 1/8 and 1/4 features (concatenation)
+        # if nshot == 1:
+        #     support_feat = support_feats[self.stack_ids[1] - 1]
+        # else:
+        support_feat = torch.stack([support_feats[k][self.stack_ids[1] - 1] for k in range(nshot)]).max(dim=0).values
+        mix = torch.cat((mix, query_feats[self.stack_ids[1] - 1], support_feat), 1)
+
+        upsample_size = (mix.size(-1) * 2,) * 2
+        mix = F.interpolate(mix, upsample_size, mode='bilinear', align_corners=True)
+        # if nshot == 1:
+        #     support_feat = support_feats[self.stack_ids[0] - 1]
+        # else:
+        support_feat = torch.stack([support_feats[k][self.stack_ids[0] - 1] for k in range(nshot)]).max(dim=0).values
+        mix = torch.cat((mix, query_feats[self.stack_ids[0] - 1], support_feat), 1)
+
+        # mixer blocks forward
+        out = self.mixer1(mix)
+        upsample_size = (out.size(-1) * 2,) * 2
+        out = F.interpolate(out, upsample_size, mode='bilinear', align_corners=True)
+        out = self.mixer2(out)
+        upsample_size = (out.size(-1) * 2,) * 2
+        out = F.interpolate(out, upsample_size, mode='bilinear', align_corners=True)
+        logit_mask = self.mixer3(out)
+
+        return logit_mask
+
+    def build_conv_block(self, in_channel, out_channels, kernel_sizes, spt_strides, group=4):
+        r""" bulid conv blocks """
+        assert len(out_channels) == len(kernel_sizes) == len(spt_strides)
+
+        building_block_layers = []
+        for idx, (outch, ksz, stride) in enumerate(zip(out_channels, kernel_sizes, spt_strides)):
+            inch = in_channel if idx == 0 else out_channels[idx - 1]
+            pad = ksz // 2
+
+            building_block_layers.append(nn.Conv2d(in_channels=inch, out_channels=outch,
+                                                   kernel_size=ksz, stride=stride, padding=pad))
+            building_block_layers.append(nn.GroupNorm(group, outch))
+            building_block_layers.append(nn.ReLU(inplace=True))
+
+        return nn.Sequential(*building_block_layers)

model/base/swin_transformer.py

@@ -0,0 +1,605 @@
+# --------------------------------------------------------
+# Swin Transformer
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ze Liu
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+
+class Mlp(nn.Module):
+    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
+
+
+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):
+    r""" 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., proj_drop=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=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """
+        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
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
+
+    def flops(self, N):
+        # calculate flops for 1 window with token length of N
+        flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        return flops
+
+
+class SwinTransformerBlock(nn.Module):
+    r""" Swin Transformer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        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, input_resolution, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        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. 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)
+
+        if self.shift_size > 0:
+            # calculate attention mask for SW-MSA
+            H, W = self.input_resolution
+            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 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))
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+
+    def forward(self, x):
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+        else:
+            shifted_x = x
+
+        # 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=self.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, H, W)  # 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
+        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
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # W-MSA/SW-MSA
+        nW = H * W / self.window_size / self.window_size
+        flops += nW * self.attn.flops(self.window_size * self.window_size)
+        # mlp
+        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
+        # norm2
+        flops += self.dim * H * W
+        return flops
+
+
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = x.view(B, H, W, C)
+
+        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
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = H * W * self.dim
+        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
+        return flops
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        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 | 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, input_resolution, depth, num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
+                                 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(input_resolution, dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        feats = []
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+            feats.append(x.clone().detach())
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return feats, x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+    def flops(self):
+        flops = 0
+        for blk in self.blocks:
+            flops += blk.flops()
+        if self.downsample is not None:
+            flops += self.downsample.flops()
+        return flops
+
+
+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, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        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.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):
+        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]})."
+        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 SwinTransformer(nn.Module):
+    r""" Swin Transformer
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+
+    Args:
+        img_size (int | tuple(int)): Input image size. Default 224
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each Swin Transformer layer.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        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. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        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
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000,
+                 embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
+                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
+                 use_checkpoint=False, feat_ids=[1, 2, 3, 4], **kwargs):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
+        self.mlp_ratio = mlp_ratio
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        # absolute position embedding
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+            trunc_normal_(self.absolute_pos_embed, std=.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),
+                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
+                                                 patches_resolution[1] // (2 ** i_layer)),
+                               depth=depths[i_layer],
+                               num_heads=num_heads[i_layer],
+                               window_size=window_size,
+                               mlp_ratio=self.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)
+
+        self.norm = norm_layer(self.num_features)
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+        self.feat_ids = feat_ids
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, 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)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        self.feat_maps = []
+        for i, layer in enumerate(self.layers):
+            feats, x = layer(x)
+            if i+1 in self.feat_ids:
+                self.feat_maps += feats
+
+        x = self.norm(x)  # B L C
+        x = self.avgpool(x.transpose(1, 2))  # B C 1
+        x = torch.flatten(x, 1)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+    def flops(self):
+        flops = 0
+        flops += self.patch_embed.flops()
+        for i, layer in enumerate(self.layers):
+            flops += layer.flops()
+        flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers)
+        flops += self.num_features * self.num_classes
+        return flops
+
+
+if __name__ == '__main__':
+    input = torch.randn(2, 3, 384, 384).cuda()
+
+    net = SwinTransformer(img_size=384, patch_size=4, window_size=12, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32))
+    net.load_state_dict(torch.load("/apdcephfs/share_1290796/shixinyu/checkpoints/swin_base_patch4_window12_384_22kto1k.pth")['model'])
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    net.to(device)
+
+    out = net.forward_features(input)
+    feat = net.feat_maps
+    for x in feat:
+        print(x.shape)

model/base/transformer.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+import math, copy
+from torch.autograd import Variable
+
+TRAIN = False
+
+class MultiHeadedAttention(nn.Module):
+    def __init__(self, h, d_model, dropout=0.1):
+        "Take in model size and number of heads."
+        super(MultiHeadedAttention, self).__init__()
+        assert d_model % h == 0
+        # We assume d_v always equals d_k
+        self.d_k = d_model // h
+        self.h = h
+        self.linears = clones(nn.Linear(d_model, d_model), 2)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, query, key, value, mask=None):
+        if mask is not None:
+            # Same mask applied to all h heads.
+            mask = mask.unsqueeze(1)
+        nbatches = query.size(0)
+        
+        # 1) Do all the linear projections in batch from d_model => h x d_k
+        query, key = \
+            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
+             for l, x in zip(self.linears, (query, key))]
+
+        value = value.repeat(self.h, 1, 1).transpose(0, 1).contiguous().unsqueeze(-1)
+
+        # query_dir, key_dir = [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) for l, x in zip([self.linears[0], self.linears[0]], (query, key))]
+        # query_norm = self.linears[1](query)[:, :, :self.h].view(nbatches, -1, self.h).transpose(1, 2)
+        # key_norm = self.linears[1](key)[:, :, :self.h].view(nbatches, -1, self.h).transpose(1, 2)
+        # query = query_dir / query_dir.norm(dim=-1).unsqueeze(-1) * 10 * query_norm.unsqueeze(-1)
+        # key = key_dir / key_dir.norm(dim=-1).unsqueeze(-1) * 10 * key_norm.unsqueeze(-1)
+        
+        if not TRAIN:
+            query = query.detach().cpu()
+            key = key.detach().cpu()
+            value = value.detach().cpu()
+        # 2) Apply attention on all the projected vectors in batch.
+        x, self.attn = attention(query, key, value, mask=mask,
+                                 dropout=self.dropout)
+        if not TRAIN:
+            x = x.cuda()
+
+        # 3) "Concat" using a view and apply a final linear.
+        return torch.mean(x, -3)
+
+
+class PositionalEncoding(nn.Module):
+    "Implement the PE function."
+
+    def __init__(self, d_model, dropout, max_len=10000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) *
+                             -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + Variable(self.pe[:, :x.size(1)],
+                         requires_grad=False)
+        return self.dropout(x)
+
+importance = torch.tensor(0.).float().cuda()
+cnt = 0
+
+def attention(query, key, value, mask=None, dropout=None):
+    "Compute 'Scaled Dot Product Attention'"
+    d_k = query.size(-1)
+    scores = torch.matmul(query, key.transpose(-2, -1)) \
+             / math.sqrt(d_k)
+    if mask is not None:
+        scores = scores.masked_fill(mask == 0, -1e9)
+    p_attn = F.softmax(scores, dim=-1)
+    # global importance, cnt
+    # im = p_attn[:, :, :, :query.size(2)].max(2)[0].mean()
+    # importance += im
+    # cnt += 1
+    if dropout is not None:
+        p_attn = dropout(p_attn)
+    return torch.matmul(p_attn, value), p_attn
+
+
+def clones(module, N):
+    "Produce N identical layers."
+    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

modelsize_estimate.py

@@ -0,0 +1,38 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def modelsize(model, input, type_size=4):
+    para = sum([np.prod(list(p.size())) for p in model.parameters()])
+    # print('Model {} : Number of params: {}'.format(model._get_name(), para))
+    print('Model {} : params: {:4f}M'.format(model._get_name(), para * type_size / 1000 / 1000))
+
+    input_ = input.clone()
+    input_.requires_grad_(requires_grad=False)
+
+    mods = list(model.modules())
+    out_sizes = []
+
+    for i in range(1, len(mods)):
+        m = mods[i]
+        if isinstance(m, nn.ReLU):
+            if m.inplace:
+                continue
+        out = m(input_)
+        out_sizes.append(np.array(out.size()))
+        input_ = out
+
+    total_nums = 0
+    for i in range(len(out_sizes)):
+        s = out_sizes[i]
+        nums = np.prod(np.array(s))
+        total_nums += nums
+
+    # print('Model {} : Number of intermedite variables without backward: {}'.format(model._get_name(), total_nums))
+    # print('Model {} : Number of intermedite variables with backward: {}'.format(model._get_name(), total_nums*2))
+    print('Model {} : intermedite variables: {:3f} M (without backward)'
+          .format(model._get_name(), total_nums * type_size / 1000 / 1000))
+    print('Model {} : intermedite variables: {:3f} M (with backward)'
+          .format(model._get_name(), total_nums * type_size*2 / 1000 / 1000))
+

scripts/importance_analysis.sh

@@ -0,0 +1,16 @@
+python ./importance_analysis.py --datapath "/research/d4/gds/wltang21/data" \
+                 --benchmark coco \
+                 --fold 0 \
+                 --bsz 1 \
+                 --nworker 1 \
+                 --backbone resnet50 \
+                 --feature_extractor_path "/research/d4/gds/wltang21/logistic_project/DCAMA/backbones/resnet50_a1h-35c100f8.pth" \
+                 --logpath "./logs" \
+		 --load "/research/d4/gds/wltang21/logistic_project/DCAMA/checkpoint/coco-20i/resnet50_fold0.pt" \
+                 --nshot 10
+
+
+		 # --load "/research/d6/rshr/xjgao/twl/logistic_project/DCAMA/checkpoint/coco-20i/resnet50_fold0.pt" \
+#                 --visualize
+#checkpoint/coco-20i/resnet50_fold0.pt
+# log/train/fold_0_ft_v0/best_model.pt

scripts/test.sh

@@ -0,0 +1,15 @@
+python ./test.py --datapath "/research/d4/gds/wltang21/data" \
+                 --benchmark coco \
+                 --fold 0 \
+                 --bsz 1 \
+                 --nworker 1 \
+                 --backbone resnet50 \
+                 --feature_extractor_path "/research/d4/gds/wltang21/logistic_project/DCAMA/backbones/resnet50_a1h-35c100f8.pth" \
+                 --logpath "./logs" \
+		 --load "/research/d4/gds/wltang21/logistic_project/DCAMA/checkpoint/coco-20i/resnet50_fold0.pt" \
+                 --nshot 30
+
+		 # --load "/research/d6/rshr/xjgao/twl/logistic_project/DCAMA/checkpoint/coco-20i/resnet50_fold0.pt" \
+#                 --visualize
+#checkpoint/coco-20i/resnet50_fold0.pt
+# log/train/fold_0_ft_v0/best_model.pt

scripts/train.sh

@@ -0,0 +1,11 @@
+python -u -m torch.distributed.launch --nnodes=1 --nproc_per_node=4 --node_rank=0 --master_port=16006 \
+./train.py --datapath "../datasets" \
+           --benchmark coco \
+           --fold 0 \
+           --bsz 12 \
+           --nworker 8 \
+           --backbone swin \
+           --feature_extractor_path "../backbones/swin_base_patch4_window12_384.pth" \
+           --logpath "./logs" \
+           --lr 1e-3 \
+           --nepoch 500

scripts/train_1gpu.sh

@@ -0,0 +1,12 @@
+python ./train_1gpu.py --datapath "/home/bkdongxianchi/MY_MOT/TWL/data" \
+           --benchmark coco \
+           --fold 0 \
+           --bsz 1 \
+           --nworker 0 \
+           --backbone resnet50 \
+           --feature_extractor_path "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/backbones/resnet50_a1h-35c100f8.pth" \
+           --logpath "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/log" \
+           --lr 1e-4 \
+           --nepoch 50 \
+           --load "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/log/resnet50_fold0.pt" \
+           --nshot 3

scripts/train_1gpu_retriver.sh

@@ -0,0 +1,12 @@
+python ./train_1gpu_retriever.py --datapath "/home/bkdongxianchi/MY_MOT/TWL/data" \
+           --benchmark coco \
+           --fold 1 \
+           --bsz 1 \
+           --nworker 0 \
+           --backbone resnet50 \
+           --feature_extractor_path "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/backbones/resnet50_a1h-35c100f8.pth" \
+           --logpath "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/log" \
+           --lr 1e-4 \
+           --nepoch 50 \
+           --load "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/log/fold_1_ft_v0/model_45.pt" \
+           --nshot 1

scripts/train_2gpu.sh

@@ -0,0 +1,14 @@
+python -u -m torch.distributed.launch --nproc_per_node=2 --master_port=18024 \
+./train.py --datapath "/home/bkdongxianchi/MY_MOT/TWL/data" \
+           --benchmark coco \
+           --fold 0 \
+           --bsz 1 \
+           --nworker 8 \
+           --backbone resnet50 \
+           --feature_extractor_path "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/backbones/resnet50_a1h-35c100f8.pth" \
+           --logpath "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/log" \
+           --lr 1e-4 \
+           --nepoch  50 \
+           --load "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/checkpoint/coco-20i/resnet50_fold0.pt" \
+           --nshot 10
+#           --load "/research/d6/rshr/xjgao/twl/logistic_project/DCAMA/checkpoint/coco-20i/resnet50_fold0.pt" \

scripts/train_2gpu_retriever.sh

@@ -0,0 +1,14 @@
+python -u -m torch.distributed.launch --nproc_per_node=2 --master_port=18024 \
+./train_retriever.py --datapath "/home/bkdongxianchi/MY_MOT/TWL/data" \
+           --benchmark coco \
+           --fold 0 \
+           --bsz 1 \
+           --nworker 8 \
+           --backbone resnet50 \
+           --feature_extractor_path "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/backbones/resnet50_a1h-35c100f8.pth" \
+           --logpath "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/log" \
+           --lr 1e-4 \
+           --nepoch  50 \
+           --load "/home/bkdongxianchi/MY_MOT/TWL/DCAMA/checkpoint/coco-20i/resnet50_fold0.pt" \
+           --nshot 1
+#           --load "/research/d6/rshr/xjgao/twl/logistic_project/DCAMA/checkpoint/coco-20i/resnet50_fold0.pt" \

scripts/train_4gpu.sh

@@ -0,0 +1,14 @@
+python -u -m torch.distributed.launch --nproc_per_node=4 --master_port=18024 \
+./train.py --datapath "~/MY_MOT/TWL/data" \
+           --benchmark coco \
+           --fold 0 \
+           --bsz 1 \
+           --nworker 8 \
+           --backbone resnet101 \
+           --feature_extractor_path "~/MY_MOT/TWL/logistic_project/DCAMA/backbones/swin_base_patch4_window12_384_22kto1k.pth" \
+           --logpath "~/MY_MOT/TWL/logistic_project/DCAMA/log" \
+           --lr 1e-4 \
+           --nepoch  50 \
+           --load "~/MY_MOT/TWL/logistic_project/DCAMA/checkpoint/coco-20i/swin_fold2.pt" \
+           --nshot 3
+#           --load "/research/d6/rshr/xjgao/twl/logistic_project/DCAMA/checkpoint/coco-20i/resnet50_fold0.pt" \

test.py

@@ -0,0 +1,132 @@
+r""" Dense Cross-Query-and-Support Attention Weighted Mask Aggregation for Few-Shot Segmentation """
+import torch.nn as nn
+import torch
+
+from model.DCAMA import DCAMA
+from common.logger import Logger, AverageMeter
+from common.vis import Visualizer
+from common.evaluation import Evaluator
+from common.config import parse_opts
+from common import utils
+from data.dataset import FSSDataset
+import cv2
+import numpy as np
+import os
+# from gpu_mem_track import MemTracker
+
+# gpu_tracker = MemTracker()
+def test(model, dataloader, nshot):
+    r""" Test """
+
+    # Freeze randomness during testing for reproducibility
+    utils.fix_randseed(0)
+    average_meter = AverageMeter(dataloader.dataset)
+
+    for idx, batch in enumerate(dataloader):
+
+        # 1. forward pass
+        nshot = batch['support_imgs'].size(1)
+        ## TODO:
+
+        batch = utils.to_cuda(batch)
+        # gpu_tracker.track()
+        pred_mask, simi, simi_map = model.module.predict_mask_nshot(batch, nshot=nshot)
+        # gpu_tracker.track()
+        torch.cuda.synchronize()
+        assert pred_mask.size() == batch['query_mask'].size()
+
+        # 2. Evaluate prediction
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask.clone(), batch)
+        
+        ## TODO:
+        iou = area_inter[1] / area_union[1]
+        
+        '''
+        cv2.imwrite('debug/query.png', cv2.imread("/home/bkdongxianchi/MY_MOT/TWL/data/COCO2014/{}".format(batch['query_name'][0])))
+        cv2.imwrite('debug/query_mask.png', (batch['query_mask'][0] * 255).detach().cpu().numpy().astype(np.uint8))
+        cv2.imwrite('debug/support_{:.3}.png'.format(iou.item()), cv2.imread('/home/bkdongxianchi/MY_MOT/TWL/data/COCO2014/{}'.format(batch['support_names'][0][0])))
+        cv2.imwrite('debug/support_mask_{:.3}.png'.format(iou.item()), (batch['support_masks'][0][0] * 255).detach().cpu().numpy().astype(np.uint8))
+        simi_map = simi_map - simi_map.min()
+        simi_map = (simi_map / simi_map.max() * 255).detach().cpu().numpy().astype(np.uint8)
+        cv2.imwrite('debug/simi_map_{:.3}.png'.format(iou.item()), simi_map)
+        
+        if os.path.exists('debug/stats.txt'):
+            with open('debug/stats.txt', "a") as f:
+                f.write("{} {}\n".format(simi.item(), iou.item()))
+        else:
+            with open('debug/stats.txt', 'w') as f:
+                f.write('{} {}\n'.format(simi.item(), iou.item()))
+        '''
+
+        average_meter.update(area_inter, area_union, batch['class_id'], loss=None)
+        average_meter.write_process(idx, len(dataloader), epoch=-1, write_batch_idx=1)
+
+        # Visualize predictions
+        if Visualizer.visualize:
+            Visualizer.visualize_prediction_batch(batch['support_imgs'], batch['support_masks'],
+                                                  batch['query_img'], batch['query_mask'],
+                                                  pred_mask, batch['class_id'], idx,
+                                                  iou_b=area_inter[1].float() / area_union[1].float())
+
+    # Write evaluation results
+    average_meter.write_result('Test', 0)
+    miou, fb_iou = average_meter.compute_iou()
+
+    return miou, fb_iou
+
+
+if __name__ == '__main__':
+
+    # Arguments parsing
+    args = parse_opts()
+
+    Logger.initialize(args, training=False)
+
+    # Model initialization
+    model = DCAMA(args.backbone, args.feature_extractor_path, args.use_original_imgsize)
+    model.eval()
+
+    # Device setup
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    Logger.info('# available GPUs: %d' % torch.cuda.device_count())
+    model = nn.DataParallel(model)
+    model.to(device)
+
+    # Load trained model
+    if args.load == '': raise Exception('Pretrained model not specified.')
+    params = model.state_dict()
+    state_dict = torch.load(args.load)
+    
+    if 'state_dict' in state_dict.keys():
+        state_dict = state_dict['state_dict']
+    state_dict2 = {}
+    for k, v in state_dict.items():
+        if 'scorer' not in k:
+            state_dict2[k] = v
+    state_dict = state_dict2
+
+    for k1, k2 in zip(list(state_dict.keys()), params.keys()):
+        state_dict[k2] = state_dict.pop(k1)
+   
+ 
+    try:
+        model.load_state_dict(state_dict, strict=True)
+    except:
+        for k in params.keys():
+            if k not in state_dict.keys():
+                state_dict[k] = params[k]
+        model.load_state_dict(state_dict, strict=True)
+
+    # Helper classes (for testing) initialization
+    Evaluator.initialize()
+    Visualizer.initialize(args.visualize, args.vispath)
+
+    # Dataset initialization
+    FSSDataset.initialize(img_size=384, datapath=args.datapath, use_original_imgsize=args.use_original_imgsize)
+    dataloader_test = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'test', args.nshot)
+
+    # Test
+    with torch.no_grad():
+        test_miou, test_fb_iou = test(model, dataloader_test, args.nshot)
+    Logger.info('Fold %d mIoU: %5.2f \t FB-IoU: %5.2f' % (args.fold, test_miou.item(), test_fb_iou.item()))
+    Logger.info('==================== Finished Testing ====================')

train.py

@@ -0,0 +1,149 @@
+r""" training (validation) code """
+import torch.optim as optim
+import torch.nn as nn
+import torch
+
+from model.DCAMA import DCAMA
+from common.logger import Logger, AverageMeter
+from common.evaluation import Evaluator
+from common.config import parse_opts
+from common import utils
+from data.dataset import FSSDataset
+import torch.nn.functional as F
+
+average_loss = torch.tensor(0.).float().cuda()
+global_idx = 0
+def train(epoch, model, dataloader, optimizer, training):
+    r""" Train """
+
+    # Force randomness during training / freeze randomness during testing
+    utils.fix_randseed(None) if training else utils.fix_randseed(0)
+    model.module.train_mode() if training else model.module.eval()
+    average_meter = AverageMeter(dataloader.dataset)
+    
+    global average_loss, global_idx
+    average_loss = average_loss.to("cuda:{}".format(torch.cuda.current_device()))
+    stats = [[], []]
+    for idx, batch in enumerate(dataloader):
+
+        # 1. forward pass
+        batch = utils.to_cuda(batch)
+        logit_mask, score_preds = model(batch['query_img'], batch['support_imgs'], batch['support_masks'], nshot=batch['support_imgs'].size(1))
+        pred_mask = logit_mask.argmax(dim=1)
+
+        # 2. Compute loss & update model parameters
+        loss = model.module.compute_objective(logit_mask, batch['query_mask'])
+        # loss_obj = loss.detach()
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask, batch)
+        iou = area_inter[1] / area_union[1]
+        loss_obj = iou.detach()
+        score_loss = F.l1_loss(score_preds, loss_obj)
+        stats[0].append(score_preds.detach().cpu().numpy())
+        stats[1].append(loss_obj.detach().cpu().numpy()[0])
+        if global_idx == 0:
+            average_loss = loss_obj.detach()
+            global_idx += 1
+        else:
+            average_loss = loss_obj.detach() * 0.05 + 0.95 * average_loss
+        print(loss_obj.item(), " ", score_preds.item(), " ", score_loss.item())
+        loss += score_loss
+
+        if training:
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        # 3. Evaluate prediction
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask, batch)
+        average_meter.update(area_inter, area_union, batch['class_id'], loss.detach().clone())
+        average_meter.write_process(idx, len(dataloader), epoch, write_batch_idx=50)
+
+    # Write evaluation results
+    average_meter.write_result('Training' if training else 'Validation', epoch)
+    avg_loss = utils.mean(average_meter.loss_buf)
+    miou, fb_iou = average_meter.compute_iou()
+    import matplotlib.pyplot as plt
+    idx = 0
+    plt.scatter(stats[0], stats[1], c="red", s=2, alpha=0.1)
+    plt.savefig('stat.png')
+    plt.close()
+
+    return avg_loss, miou, fb_iou
+
+
+if __name__ == '__main__':
+
+    # Arguments parsing
+    args = parse_opts()
+
+    # ddp backend initialization
+    torch.distributed.init_process_group(backend='nccl')
+    torch.cuda.set_device(args.local_rank)
+
+    # Model initialization
+    model = DCAMA(args.backbone, args.feature_extractor_path, False)
+    device = torch.device("cuda", args.local_rank)
+    model.to(device)
+    model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank,
+                                                find_unused_parameters=True)
+    
+
+    params = model.state_dict()
+    state_dict = torch.load(args.load)
+    state_dict2 = {}
+    for k in state_dict.keys():
+        if "scorer" in k:
+            continue
+        state_dict2[k] = state_dict[k]
+    state_dict = state_dict2
+    for k1, k2 in zip(list(state_dict.keys()), params.keys()):
+        state_dict[k2] = state_dict.pop(k1)
+    
+    model.load_state_dict(state_dict, strict=False)
+    
+    ## TODO:
+    for i in range(len(model.module.model.DCAMA_blocks)):
+        torch.nn.init.constant_(model.module.model.DCAMA_blocks[i].linears[1].weight, 0.)
+        torch.nn.init.constant_(model.module.model.DCAMA_blocks[i].linears[1].bias, 1.)
+
+
+    # Helper classes (for training) initialization
+    optimizer = optim.SGD([{"params": model.module.model.parameters(), "lr": args.lr,
+                            "momentum": 0.9, "weight_decay": args.lr/10, "nesterov": True}])
+    Evaluator.initialize()
+    if args.local_rank == 0:
+        Logger.initialize(args, training=True)
+        Logger.info('# available GPUs: %d' % torch.cuda.device_count())
+
+    # Dataset initialization
+    FSSDataset.initialize(img_size=384, datapath=args.datapath, use_original_imgsize=False)
+    dataloader_trn = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'trn', args.nshot)
+    if args.local_rank == 0:
+        dataloader_val = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'val', args.nshot)
+
+    # Train
+    best_val_miou = float('-inf')
+    best_val_loss = float('inf')
+    for epoch in range(args.nepoch):
+        dataloader_trn.sampler.set_epoch(epoch)
+        trn_loss, trn_miou, trn_fb_iou = train(epoch, model, dataloader_trn, optimizer, training=True)
+
+        # evaluation
+        if args.local_rank == 0:
+            # with torch.no_grad():
+            #     val_loss, val_miou, val_fb_iou = train(epoch, model, dataloader_val, optimizer, training=False)
+
+            # Save the best model
+            # if val_miou > best_val_miou:
+            #     best_val_miou = val_miou
+            #     Logger.save_model_miou(model, epoch, val_miou)
+            Logger.save_model_miou(model, epoch , 1.)
+
+            # Logger.tbd_writer.add_scalars('data/loss', {'trn_loss': trn_loss, 'val_loss': val_loss}, epoch)
+            # Logger.tbd_writer.add_scalars('data/miou', {'trn_miou': trn_miou, 'val_miou': val_miou}, epoch)
+            # Logger.tbd_writer.add_scalars('data/fb_iou', {'trn_fb_iou': trn_fb_iou, 'val_fb_iou': val_fb_iou}, epoch)
+            # Logger.tbd_writer.flush()
+
+    if args.local_rank == 0:
+        Logger.tbd_writer.close()
+        Logger.info('==================== Finished Training ====================')

train_1gpu.py

@@ -0,0 +1,170 @@
+r""" training (validation) code """
+import torch.optim as optim
+import torch.nn as nn
+import torch
+
+from model.DCAMA import DCAMA
+from common.logger import Logger, AverageMeter
+from common.evaluation import Evaluator
+from common.config import parse_opts
+from common import utils
+from data.dataset import FSSDataset # FSDataset4SAM
+# from transformers import SamProcessor
+from PIL import Image
+import numpy as np
+import torch.nn.functional as F
+from torchvision import transforms
+import pickle
+import pycocotools.coco as COCO
+import cv2
+
+def train(epoch, model, dataloader, optimizer, training, shot=1):
+    r""" Train """
+
+    # Force randomness during training / freeze randomness during testing
+    utils.fix_randseed(None) if training else utils.fix_randseed(0)
+
+    if hasattr(model, "module"):
+        model.module.train_mode() if training else model.module.eval()
+    else:
+        model.train_mode() if training else model.module.eval()
+    average_meter = AverageMeter(dataloader.dataset)
+    average_loss = torch.tensor(0.).float().cuda()
+    stats = [[], []]
+    criterion_score = nn.BCEWithLogitsLoss()
+    for idx, batch in enumerate(dataloader):
+        
+        # batch = process_batch4SAM(batch)
+        shot = batch['support_imgs'].size(1)
+        # 1. forward pass
+        batch = utils.to_cuda(batch)
+        logit_mask, score_preds = model(batch['query_img'], batch['support_imgs'], batch['support_masks'], nshot=shot)
+        pred_mask = logit_mask.argmax(dim=1)
+        # 2. Compute loss & update model parameters
+        loss = model.compute_objective(logit_mask, batch['query_mask'])
+        # loss_obj = loss.detach()
+        
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask, batch)
+        
+        iou = (area_inter[1] / area_union[1]).float()
+        if iou > 0.7 or iou  < 0.1:
+            '''
+            if iou < 0.1:
+                img = batch['query_img'][0].permute(1, 2, 0).detach().cpu().numpy()
+                img = img - img.min()
+                img = img / img.max()
+                cv2.imwrite('query_image.png', (img * 255).astype(np.uint8))
+                img = batch['support_imgs'][0][0].permute(1, 2, 0).detach().cpu().numpy()
+                img = img - img.min()
+                img = img / img.max()
+                cv2.imwrite('support_image.png', (img * 255).astype(np.uint8))
+                cv2.imwrite('query_mask.png', (batch['query_mask'][0] * 255).detach().cpu().numpy().astype(np.uint8))
+                cv2.imwrite('pred_mask.png', (pred_mask[0] * 255).detach().cpu().numpy().astype(np.uint8))
+                cv2.imwrite('support_mask.png', (batch['support_masks'][0][0] * 255).detach().cpu().numpy().astype(np.uint8))
+            '''
+            if iou > 0.7:
+                iou = torch.tensor(1.).float().cuda()
+            else:
+                iou = torch.tensor(0.).float().cuda()
+            score_loss = criterion_score(score_preds, iou)
+            stats[0].append(score_preds.detach().cpu().numpy())
+            stats[1].append((area_inter[1] / area_union[1]).detach().cpu().numpy())
+            print(score_preds, (area_inter[1] / area_union[1]))
+            
+
+        if training:
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+        # 3. Evaluate prediction
+
+        # img = batch['support_imgs'][0][0].permute(1, 2, 0)
+        # img = img - img.min()
+        # img /= img.max()
+        # import cv2
+        # cv2.imwrite("debug.png", (img * 255).detach().cpu().numpy())
+        # cv2.imwrite("debug2.png", (batch['support_masks'][0][0] * 255).detach().cpu().numpy())
+        # import ipdb;ipdb.set_trace()
+
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask, batch)
+        average_meter.update(area_inter, area_union, batch['class_id'], loss.detach().clone())
+        average_meter.write_process(idx, len(dataloader), epoch, write_batch_idx=50)
+
+    # Write evaluation results
+    average_meter.write_result('Training' if training else 'Validation', epoch)
+    avg_loss = utils.mean(average_meter.loss_buf)
+    miou, fb_iou = average_meter.compute_iou()
+    
+    import matplotlib.pyplot as plt
+    plt.scatter(stats[0], stats[1], c="red", s=2, alpha=0.02)
+    plt.savefig("stats.png")
+    return avg_loss, miou, fb_iou
+
+
+if __name__ == '__main__':
+
+    # Arguments parsing
+    args = parse_opts()
+
+    # Model initialization
+    model = DCAMA(args.backbone, args.feature_extractor_path, False)
+    device = torch.device("cuda", args.local_rank)
+    model.to(device)
+
+    params = model.state_dict()
+    state_dict = torch.load(args.load)
+    if 'state_dict' in state_dict.keys():
+        state_dict = state_dict['state_dict']
+    state_dict2 = {}
+    for k in state_dict.keys():
+        if "scorer" in k:
+            continue
+        state_dict2[k] = state_dict[k]
+    state_dict = state_dict2
+    for k1, k2 in zip(list(state_dict.keys()), params.keys()):
+        state_dict[k2] = state_dict.pop(k1)
+    model.load_state_dict(state_dict, strict=False)
+    
+    ## TODO:
+    for i in range(len(model.model.DCAMA_blocks)):
+        torch.nn.init.constant_(model.model.DCAMA_blocks[i].linears[1].weight, 0.)
+        torch.nn.init.constant_(model.model.DCAMA_blocks[i].linears[1].bias, 1.)
+    # Helper classes (for training) initialization
+    optimizer = optim.SGD([{"params": model.parameters(), "lr": args.lr,
+                            "momentum": 0.9, "weight_decay": args.lr/10, "nesterov": True}])
+    Evaluator.initialize()
+    if args.local_rank == 0:
+        Logger.initialize(args, training=True)
+        Logger.info('# available GPUs: %d' % torch.cuda.device_count())
+
+    # Dataset initialization
+    FSSDataset.initialize(img_size=384, datapath=args.datapath, use_original_imgsize=False)
+    dataloader_trn = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'trn', shot=args.nshot)
+    if args.local_rank == 0:
+        dataloader_val = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'val', shot=args.nshot)
+
+    # Train
+    best_val_miou = float('-inf')
+    best_val_loss = float('inf')
+
+    for epoch in range(args.nepoch):
+        trn_loss, trn_miou, trn_fb_iou = train(epoch, model, dataloader_trn, optimizer, training=True, shot=args.nshot)
+
+        # evaluation
+        if args.local_rank == 0:
+            # with torch.no_grad():
+            #     val_loss, val_miou, val_fb_iou = train(epoch, model, dataloader_val, optimizer, training=False)
+
+            # Save the best model
+            # if val_miou > best_val_miou:
+            #     best_val_miou = val_miou
+            Logger.save_model_miou(model, epoch, 1.)
+
+            # Logger.tbd_writer.add_scalars('data/loss', {'trn_loss': trn_loss, 'val_loss': val_loss}, epoch)
+            # Logger.tbd_writer.add_scalars('data/miou', {'trn_miou': trn_miou, 'val_miou': val_miou}, epoch)
+            # Logger.tbd_writer.add_scalars('data/fb_iou', {'trn_fb_iou': trn_fb_iou, 'val_fb_iou': val_fb_iou}, epoch)
+            # Logger.tbd_writer.flush()
+
+    if args.local_rank == 0:
+        Logger.tbd_writer.close()
+        Logger.info('==================== Finished Training ====================')

train_1gpu_retriever.py

@@ -0,0 +1,172 @@
+r""" training (validation) code """
+import torch.optim as optim
+import torch.nn as nn
+import torch
+
+from model.DCAMA import DCAMA
+from common.logger import Logger, AverageMeter
+from common.evaluation import Evaluator
+from common.config import parse_opts
+from common import utils
+from data.dataset import FSSDataset # FSDataset4SAM
+# from transformers import SamProcessor
+from PIL import Image
+import numpy as np
+import torch.nn.functional as F
+from torchvision import transforms
+import pickle
+import pycocotools.coco as COCO
+import cv2
+import torchvision
+
+def train(epoch, model, dataloader, optimizer, training, shot=1):
+    r""" Train """
+
+    # Force randomness during training / freeze randomness during testing
+    utils.fix_randseed(None) if training else utils.fix_randseed(0)
+
+    if hasattr(model, "module"):
+        model.module.train_mode() if training else model.module.eval()
+    else:
+        model.train_mode() if training else model.module.eval()
+    average_meter = AverageMeter(dataloader.dataset)
+    average_loss = torch.tensor(0.).float().cuda()
+    stats = [[], []]
+    criterion_score = nn.BCEWithLogitsLoss()
+    for idx, batch in enumerate(dataloader):
+        
+        # batch = process_batch4SAM(batch)
+        shot = batch['support_imgs'].size(1)
+        # 1. forward pass
+        batch = utils.to_cuda(batch)
+        logit_mask, score_preds = model(batch['query_img'], batch['support_imgs'], batch['support_masks'], nshot=shot, predict_score=True)
+        pred_mask = logit_mask.argmax(dim=1)
+        # 2. Compute loss & update model parameters
+        loss = model.compute_objective(logit_mask, batch['query_mask'])
+        # loss_obj = loss.detach()
+        
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask, batch)
+        
+        iou = (area_inter[1] / area_union[1]).float()
+        if iou > 0.7 or iou < 0.05:
+            '''
+            if iou < 0.1:
+                img = batch['query_img'][0].permute(1, 2, 0).detach().cpu().numpy()
+                img = img - img.min()
+                img = img / img.max()
+                cv2.imwrite('query_image.png', (img * 255).astype(np.uint8))
+                img = batch['support_imgs'][0][0].permute(1, 2, 0).detach().cpu().numpy()
+                img = img - img.min()
+                img = img / img.max()
+                cv2.imwrite('support_image.png', (img * 255).astype(np.uint8))
+                cv2.imwrite('query_mask.png', (batch['query_mask'][0] * 255).detach().cpu().numpy().astype(np.uint8))
+                cv2.imwrite('pred_mask.png', (pred_mask[0] * 255).detach().cpu().numpy().astype(np.uint8))
+                cv2.imwrite('support_mask.png', (batch['support_masks'][0][0] * 255).detach().cpu().numpy().astype(np.uint8))
+            '''
+            if iou > 0.7:
+                iou = torch.tensor(1.).float().cuda()
+            else:
+                iou = torch.tensor(0.).float().cuda()
+            score_loss = torchvision.ops.sigmoid_focal_loss(score_preds, iou)
+            # score_loss = F.l1_loss(score_preds, iou)
+            stats[0].append(score_preds.detach().cpu().numpy())
+            stats[1].append((area_inter[1] / area_union[1]).detach().cpu().numpy())
+            print(score_preds, (area_inter[1] / area_union[1]))
+            loss = score_loss
+
+            if training:
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+        # 3. Evaluate prediction
+
+        # img = batch['support_imgs'][0][0].permute(1, 2, 0)
+        # img = img - img.min()
+        # img /= img.max()
+        # import cv2
+        # cv2.imwrite("debug.png", (img * 255).detach().cpu().numpy())
+        # cv2.imwrite("debug2.png", (batch['support_masks'][0][0] * 255).detach().cpu().numpy())
+        # import ipdb;ipdb.set_trace()
+
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask, batch)
+        average_meter.update(area_inter, area_union, batch['class_id'], loss.detach().clone())
+        average_meter.write_process(idx, len(dataloader), epoch, write_batch_idx=50)
+
+    # Write evaluation results
+    average_meter.write_result('Training' if training else 'Validation', epoch)
+    avg_loss = utils.mean(average_meter.loss_buf)
+    miou, fb_iou = average_meter.compute_iou()
+    
+    import matplotlib.pyplot as plt
+    plt.scatter(stats[0], stats[1], c="red", s=2, alpha=0.02)
+    plt.savefig("stats.png")
+    return avg_loss, miou, fb_iou
+
+
+if __name__ == '__main__':
+
+    # Arguments parsing
+    args = parse_opts()
+
+    # Model initialization
+    model = DCAMA(args.backbone, args.feature_extractor_path, False)
+    device = torch.device("cuda", args.local_rank)
+    model.to(device)
+
+    params = model.state_dict()
+    state_dict = torch.load(args.load)
+    if 'state_dict' in state_dict.keys():
+        state_dict = state_dict['state_dict']
+    state_dict2 = {}
+    for k in state_dict.keys():
+        if "scorer" in k:
+            continue
+        state_dict2[k] = state_dict[k]
+    state_dict = state_dict2
+    for k1, k2 in zip(list(state_dict.keys()), params.keys()):
+        state_dict[k2] = state_dict.pop(k1)
+    model.load_state_dict(state_dict, strict=False)
+    
+    ## TODO:
+    # for i in range(len(model.model.DCAMA_blocks)):
+    #     torch.nn.init.constant_(model.model.DCAMA_blocks[i].linears[1].weight, 0.)
+    #     torch.nn.init.constant_(model.model.DCAMA_blocks[i].linears[1].bias, 1.)
+    # Helper classes (for training) initialization
+    optimizer = optim.SGD([{"params": model.parameters(), "lr": args.lr,
+                            "momentum": 0.9, "weight_decay": args.lr/10, "nesterov": True}])
+    Evaluator.initialize()
+    if args.local_rank == 0:
+        Logger.initialize(args, training=True)
+        Logger.info('# available GPUs: %d' % torch.cuda.device_count())
+
+    # Dataset initialization
+    FSSDataset.initialize(img_size=384, datapath=args.datapath, use_original_imgsize=False)
+    dataloader_trn = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'trn', shot=args.nshot)
+    if args.local_rank == 0:
+        dataloader_val = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'val', shot=args.nshot)
+
+    # Train
+    best_val_miou = float('-inf')
+    best_val_loss = float('inf')
+
+    for epoch in range(args.nepoch):
+        trn_loss, trn_miou, trn_fb_iou = train(epoch, model, dataloader_trn, optimizer, training=True, shot=args.nshot)
+
+        # evaluation
+        if args.local_rank == 0:
+            # with torch.no_grad():
+            #     val_loss, val_miou, val_fb_iou = train(epoch, model, dataloader_val, optimizer, training=False)
+
+            # Save the best model
+            # if val_miou > best_val_miou:
+            #     best_val_miou = val_miou
+            Logger.save_model_miou(model, epoch, 1.)
+
+            # Logger.tbd_writer.add_scalars('data/loss', {'trn_loss': trn_loss, 'val_loss': val_loss}, epoch)
+            # Logger.tbd_writer.add_scalars('data/miou', {'trn_miou': trn_miou, 'val_miou': val_miou}, epoch)
+            # Logger.tbd_writer.add_scalars('data/fb_iou', {'trn_fb_iou': trn_fb_iou, 'val_fb_iou': val_fb_iou}, epoch)
+            # Logger.tbd_writer.flush()
+
+    if args.local_rank == 0:
+        Logger.tbd_writer.close()
+        Logger.info('==================== Finished Training ====================')

train_retriever.py

@@ -0,0 +1,164 @@
+r""" training (validation) code """
+import torch.optim as optim
+import torch.nn as nn
+import torch
+
+from model.DCAMA import DCAMA
+from common.logger import Logger, AverageMeter
+from common.evaluation import Evaluator
+from common.config import parse_opts
+from common import utils
+from data.dataset import FSSDataset
+import torch.nn.functional as F
+
+average_loss = torch.tensor(0.).float().cuda()
+global_idx = 0
+def train(epoch, model, dataloader, optimizer, training):
+    r""" Train """
+
+    # Force randomness during training / freeze randomness during testing
+    utils.fix_randseed(None) if training else utils.fix_randseed(0)
+    model.module.train_mode() if training else model.module.eval()
+    average_meter = AverageMeter(dataloader.dataset)
+    
+    global average_loss, global_idx
+    average_loss = average_loss.to("cuda:{}".format(torch.cuda.current_device()))
+    stats = [[], []]
+    criterion_score = nn.BCEWithLogitsLoss()
+    for idx, batch in enumerate(dataloader):
+
+        # 1. forward pass
+        batch = utils.to_cuda(batch)
+        logit_mask, score_preds = model(batch['query_img'], batch['support_imgs'], batch['support_masks'], nshot=batch['support_imgs'].size(1))
+        pred_mask = logit_mask.argmax(dim=1)
+
+        # 2. Compute loss & update model parameters
+        loss = model.module.compute_objective(logit_mask, batch['query_mask'])
+        # loss_obj = loss.detach()
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask, batch)
+        iou = area_inter[1] / area_union[1]
+        
+        if iou > 0.7 or iou  < 0.1:
+            '''
+            if iou < 0.1:
+                img = batch['query_img'][0].permute(1, 2, 0).detach().cpu().numpy()
+                img = img - img.min()
+                img = img / img.max()
+                cv2.imwrite('query_image.png', (img * 255).astype(np.uint8))
+                img = batch['support_imgs'][0][0].permute(1, 2, 0).detach().cpu().numpy()
+                img = img - img.min()
+                img = img / img.max()
+                cv2.imwrite('support_image.png', (img * 255).astype(np.uint8))
+                cv2.imwrite('query_mask.png', (batch['query_mask'][0] * 255).detach().cpu().numpy().astype(np.uint8))
+                cv2.imwrite('pred_mask.png', (pred_mask[0] * 255).detach().cpu().numpy().astype(np.uint8))
+                cv2.imwrite('support_mask.png', (batch['support_masks'][0][0] * 255).detach().cpu().numpy().astype(np.uint8))
+            '''
+            if iou > 0.7:
+                iou = torch.tensor(1.).float().cuda()
+            else:
+                iou = torch.tensor(0.).float().cuda()
+            score_loss = criterion_score(score_preds, iou)
+            stats[0].append(score_preds.detach().cpu().numpy())
+            stats[1].append((area_inter[1] / area_union[1]).detach().cpu().numpy())
+            print(score_preds, (area_inter[1] / area_union[1]))
+            loss = score_loss
+
+        if training:
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        # 3. Evaluate prediction
+        area_inter, area_union = Evaluator.classify_prediction(pred_mask, batch)
+        average_meter.update(area_inter, area_union, batch['class_id'], loss.detach().clone())
+        average_meter.write_process(idx, len(dataloader), epoch, write_batch_idx=50)
+
+    # Write evaluation results
+    average_meter.write_result('Training' if training else 'Validation', epoch)
+    avg_loss = utils.mean(average_meter.loss_buf)
+    miou, fb_iou = average_meter.compute_iou()
+    import matplotlib.pyplot as plt
+    idx = 0
+    plt.scatter(stats[0], stats[1], c="red", s=2, alpha=0.1)
+    plt.savefig('stat.png')
+    plt.close()
+
+    return avg_loss, miou, fb_iou
+
+
+if __name__ == '__main__':
+
+    # Arguments parsing
+    args = parse_opts()
+
+    # ddp backend initialization
+    torch.distributed.init_process_group(backend='nccl')
+    torch.cuda.set_device(args.local_rank)
+
+    # Model initialization
+    model = DCAMA(args.backbone, args.feature_extractor_path, False)
+    device = torch.device("cuda", args.local_rank)
+    model.to(device)
+    model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank,
+                                                find_unused_parameters=True)
+    
+
+    params = model.state_dict()
+    state_dict = torch.load(args.load)
+    state_dict2 = {}
+    for k in state_dict.keys():
+        if "scorer" in k:
+            continue
+        state_dict2[k] = state_dict[k]
+    state_dict = state_dict2
+    for k1, k2 in zip(list(state_dict.keys()), params.keys()):
+        state_dict[k2] = state_dict.pop(k1)
+    
+    model.load_state_dict(state_dict, strict=False)
+    
+    ## TODO:
+    for i in range(len(model.module.model.DCAMA_blocks)):
+        torch.nn.init.constant_(model.module.model.DCAMA_blocks[i].linears[1].weight, 0.)
+        torch.nn.init.constant_(model.module.model.DCAMA_blocks[i].linears[1].bias, 1.)
+
+
+    # Helper classes (for training) initialization
+    optimizer = optim.SGD([{"params": model.module.model.parameters(), "lr": args.lr,
+                            "momentum": 0.9, "weight_decay": args.lr/10, "nesterov": True}])
+    Evaluator.initialize()
+    if args.local_rank == 0:
+        Logger.initialize(args, training=True)
+        Logger.info('# available GPUs: %d' % torch.cuda.device_count())
+
+    # Dataset initialization
+    FSSDataset.initialize(img_size=384, datapath=args.datapath, use_original_imgsize=False)
+    dataloader_trn = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'trn', args.nshot)
+    if args.local_rank == 0:
+        dataloader_val = FSSDataset.build_dataloader(args.benchmark, args.bsz, args.nworker, args.fold, 'val', args.nshot)
+
+    # Train
+    best_val_miou = float('-inf')
+    best_val_loss = float('inf')
+    for epoch in range(args.nepoch):
+        dataloader_trn.sampler.set_epoch(epoch)
+        trn_loss, trn_miou, trn_fb_iou = train(epoch, model, dataloader_trn, optimizer, training=True)
+
+        # evaluation
+        if args.local_rank == 0:
+            # with torch.no_grad():
+            #     val_loss, val_miou, val_fb_iou = train(epoch, model, dataloader_val, optimizer, training=False)
+
+            # Save the best model
+            # if val_miou > best_val_miou:
+            #     best_val_miou = val_miou
+            #     Logger.save_model_miou(model, epoch, val_miou)
+            Logger.save_model_miou(model, epoch , 1.)
+
+            # Logger.tbd_writer.add_scalars('data/loss', {'trn_loss': trn_loss, 'val_loss': val_loss}, epoch)
+            # Logger.tbd_writer.add_scalars('data/miou', {'trn_miou': trn_miou, 'val_miou': val_miou}, epoch)
+            # Logger.tbd_writer.add_scalars('data/fb_iou', {'trn_fb_iou': trn_fb_iou, 'val_fb_iou': val_fb_iou}, epoch)
+            # Logger.tbd_writer.flush()
+
+    if args.local_rank == 0:
+        Logger.tbd_writer.close()
+        Logger.info('==================== Finished Training ====================')