Upload 127 files

CAGE_expression_inference-apvit/apvit_mmcls/__init__.py

@@ -0,0 +1,3 @@
+from .version import __version__
+
+__all__ = ['__version__']

CAGE_expression_inference-apvit/apvit_mmcls/apis/__init__.py

@@ -0,0 +1,8 @@
+from .inference import inference_model, init_model, show_result_pyplot
+from .test import multi_gpu_test, single_gpu_test
+from .train import set_random_seed, train_model
+
+__all__ = [
+    'set_random_seed', 'train_model', 'init_model', 'inference_model',
+    'multi_gpu_test', 'single_gpu_test', 'show_result_pyplot'
+]

CAGE_expression_inference-apvit/apvit_mmcls/apis/inference.py

@@ -0,0 +1,108 @@
+import warnings
+
+import matplotlib.pyplot as plt
+import mmcv
+import numpy as np
+import torch
+from mmcv.parallel import collate, scatter
+from mmcv.runner import load_checkpoint
+
+from mmcls.datasets.pipelines import Compose
+from mmcls.models import build_classifier
+
+
+def init_model(config, checkpoint=None, device='cuda:0', options=None):
+    """Initialize a classifier from config file.
+
+    Args:
+        config (str or :obj:`mmcv.Config`): Config file path or the config
+            object.
+        checkpoint (str, optional): Checkpoint path. If left as None, the model
+            will not load any weights.
+        options (dict): Options to override some settings in the used config.
+
+    Returns:
+        nn.Module: The constructed classifier.
+    """
+    if isinstance(config, str):
+        config = mmcv.Config.fromfile(config)
+    elif not isinstance(config, mmcv.Config):
+        raise TypeError('config must be a filename or Config object, '
+                        f'but got {type(config)}')
+    if options is not None:
+        config.merge_from_dict(options)
+    config.model.pretrained = None
+    config.model.extractor.pretrained = None
+    config.model.vit.pretrained = None
+    model = build_classifier(config.model)
+    if checkpoint is not None:
+        map_loc = 'cpu' if device == 'cpu' else None
+        checkpoint = load_checkpoint(model, checkpoint, map_location=map_loc)
+        class_loaded = False
+        if 'meta' in checkpoint:
+            if 'CLASSES' in checkpoint['meta']:
+                model.CLASSES = checkpoint['meta']['CLASSES']
+                class_loaded = True
+        if not class_loaded:
+            from mmcls.datasets.raf import FER_CLASSES
+            model.CLASSES = FER_CLASSES
+    model.cfg = config  # save the config in the model for convenience
+    model.to(device)
+    model.eval()
+    return model
+
+
+def inference_model(model, img):
+    """Inference image(s) with the classifier.
+
+    Args:
+        model (nn.Module): The loaded classifier.
+        img (str/ndarray): The image filename or loaded image.
+
+    Returns:
+        result (dict): The classification results that contains
+            `class_name`, `pred_label` and `pred_score`.
+    """
+    cfg = model.cfg
+    device = next(model.parameters()).device  # model device
+    # build the data pipeline
+    if isinstance(img, str):
+        if cfg.data.test.pipeline[0]['type'] != 'LoadImageFromFile':
+            cfg.data.test.pipeline.insert(0, dict(type='LoadImageFromFile'))
+        data = dict(img_info=dict(filename=img), img_prefix=None)
+    else:
+        if cfg.data.test.pipeline[0]['type'] == 'LoadImageFromFile':
+            cfg.data.test.pipeline.pop(0)
+        data = dict(img=img)
+    test_pipeline = Compose(cfg.data.test.pipeline)
+    data = test_pipeline(data)
+    data = collate([data], samples_per_gpu=1)
+    if next(model.parameters()).is_cuda:
+        # scatter to specified GPU
+        data = scatter(data, [device])[0]
+
+    # forward the model
+    with torch.no_grad():
+        scores = model(return_loss=False, **data)
+        pred_score = np.max(scores, axis=1)[0]
+        pred_label = np.argmax(scores, axis=1)[0]
+        result = {'pred_label': pred_label, 'pred_score': float(pred_score)}
+    result['pred_class'] = model.CLASSES[result['pred_label']]
+    return result
+
+
+def show_result_pyplot(model, img, result, fig_size=(15, 10)):
+    """Visualize the classification results on the image.
+
+    Args:
+        model (nn.Module): The loaded classifier.
+        img (str or np.ndarray): Image filename or loaded image.
+        result (list): The classification result.
+        fig_size (tuple): Figure size of the pyplot figure.
+    """
+    if hasattr(model, 'module'):
+        model = model.module
+    img = model.show_result(img, result, show=False)
+    plt.figure(figsize=fig_size)
+    plt.imshow(mmcv.bgr2rgb(img))
+    plt.show()

CAGE_expression_inference-apvit/apvit_mmcls/apis/test.py

@@ -0,0 +1,158 @@
+import os.path as osp
+import pickle
+import shutil
+import tempfile
+import time
+
+import mmcv
+import torch
+import torch.distributed as dist
+from mmcv.runner import get_dist_info
+
+
+def single_gpu_test(model, data_loader, show=False, out_dir=None):
+    model.eval()
+    results = []
+    dataset = data_loader.dataset
+    prog_bar = mmcv.ProgressBar(len(dataset))
+    for i, data in enumerate(data_loader):
+        with torch.no_grad():
+            result = model(return_loss=False, **data)
+        results.append(result)
+
+        if show or out_dir:
+            pass  # TODO
+
+        batch_size = data['img'].size(0)
+        for _ in range(batch_size):
+            prog_bar.update()
+    return results
+
+
+def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False):
+    """Test model with multiple gpus.
+
+    This method tests model with multiple gpus and collects the results
+    under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'
+    it encodes results to gpu tensors and use gpu communication for results
+    collection. On cpu mode it saves the results on different gpus to 'tmpdir'
+    and collects them by the rank 0 worker.
+
+    Args:
+        model (nn.Module): Model to be tested.
+        data_loader (nn.Dataloader): Pytorch data loader.
+        tmpdir (str): Path of directory to save the temporary results from
+            different gpus under cpu mode.
+        gpu_collect (bool): Option to use either gpu or cpu to collect results.
+
+    Returns:
+        list: The prediction results.
+    """
+    model.eval()
+    results = []
+    dataset = data_loader.dataset
+    rank, world_size = get_dist_info()
+    if rank == 0:
+        prog_bar = mmcv.ProgressBar(len(dataset))
+    time.sleep(2)  # This line can prevent deadlock problem in some cases.
+    for i, data in enumerate(data_loader):
+        with torch.no_grad():
+            result = model(return_loss=False, **data)
+        if isinstance(result, list):
+            results.extend(result)
+        else:
+            results.append(result)
+
+        # if rank == 0:
+        #     batch_size = data['img'].size(0)
+        #     for _ in range(batch_size * world_size):
+        #         prog_bar.update()
+
+    # collect results from all ranks
+    if gpu_collect:
+        results = collect_results_gpu(results, len(dataset))
+    else:
+        results = collect_results_cpu(results, len(dataset), tmpdir)
+    return results
+
+
+def collect_results_cpu(result_part, size, tmpdir=None):
+    rank, world_size = get_dist_info()
+    # create a tmp dir if it is not specified
+    if tmpdir is None:
+        MAX_LEN = 512
+        # 32 is whitespace
+        dir_tensor = torch.full((MAX_LEN, ),
+                                32,
+                                dtype=torch.uint8,
+                                device='cuda')
+        if rank == 0:
+            tmpdir = tempfile.mkdtemp()
+            tmpdir = torch.tensor(
+                bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda')
+            dir_tensor[:len(tmpdir)] = tmpdir
+        dist.broadcast(dir_tensor, 0)
+        tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip()
+    else:
+        mmcv.mkdir_or_exist(tmpdir)
+    # dump the part result to the dir
+    mmcv.dump(result_part, osp.join(tmpdir, f'part_{rank}.pkl'))
+    dist.barrier()
+    # collect all parts
+    if rank != 0:
+        return None
+    else:
+        # load results of all parts from tmp dir
+        part_list = []
+        for i in range(world_size):
+            part_file = osp.join(tmpdir, f'part_{i}.pkl')
+            part_result = mmcv.load(part_file)
+            # When data is severely insufficient, an empty part_result
+            # on a certain gpu could makes the overall outputs empty.
+            if part_result:
+                part_list.append(part_result)
+        # sort the results
+        ordered_results = []
+        for res in zip(*part_list):
+            ordered_results.extend(list(res))
+        # the dataloader may pad some samples
+        ordered_results = ordered_results[:size]
+        # remove tmp dir
+        shutil.rmtree(tmpdir)
+        return ordered_results
+
+
+def collect_results_gpu(result_part, size):
+    rank, world_size = get_dist_info()
+    # dump result part to tensor with pickle
+    part_tensor = torch.tensor(
+        bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')
+    # gather all result part tensor shape
+    shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
+    shape_list = [shape_tensor.clone() for _ in range(world_size)]
+    dist.all_gather(shape_list, shape_tensor)
+    # padding result part tensor to max length
+    shape_max = torch.tensor(shape_list).max()
+    part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')
+    part_send[:shape_tensor[0]] = part_tensor
+    part_recv_list = [
+        part_tensor.new_zeros(shape_max) for _ in range(world_size)
+    ]
+    # gather all result part
+    dist.all_gather(part_recv_list, part_send)
+
+    if rank == 0:
+        part_list = []
+        for recv, shape in zip(part_recv_list, shape_list):
+            part_result = pickle.loads(recv[:shape[0]].cpu().numpy().tobytes())
+            # When data is severely insufficient, an empty part_result
+            # on a certain gpu could makes the overall outputs empty.
+            if part_result:
+                part_list.append(part_result)
+        # sort the results
+        ordered_results = []
+        for res in zip(*part_list):
+            ordered_results.extend(list(res))
+        # the dataloader may pad some samples
+        ordered_results = ordered_results[:size]
+        return ordered_results

CAGE_expression_inference-apvit/apvit_mmcls/apis/train.py

@@ -0,0 +1,140 @@
+import random
+import warnings
+
+import numpy as np
+import torch
+from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
+from mmcv.runner import DistSamplerSeedHook, build_optimizer, build_runner, Fp16OptimizerHook, OptimizerHook
+
+from mmcls.core import (DistEvalHook, EvalHook)
+from mmcls.datasets import build_dataloader, build_dataset
+from mmcls.utils import get_root_logger
+
+
+def set_random_seed(seed, deterministic=False):
+    """Set random seed.
+
+    Args:
+        seed (int): Seed to be used.
+        deterministic (bool): Whether to set the deterministic option for
+            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`
+            to True and `torch.backends.cudnn.benchmark` to False.
+            Default: False.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+def train_model(model,
+                dataset,
+                cfg,
+                distributed=False,
+                validate=False,
+                timestamp=None,
+                meta=None):
+    logger = get_root_logger(cfg.log_level)
+
+    # prepare data loaders
+    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
+
+    data_loaders = [
+        build_dataloader(
+            ds,
+            cfg.data.samples_per_gpu,
+            cfg.data.workers_per_gpu,
+            # cfg.gpus will be ignored if distributed
+            num_gpus=len(cfg.gpu_ids),
+            dist=distributed,
+            round_up=True,
+            seed=cfg.seed,
+            # persistent_workers=True
+        ) for ds in dataset
+    ]
+
+    # build runner
+    optimizer = build_optimizer(model, cfg.optimizer)
+
+    # if 'use_fp16' in cfg and cfg.use_fp16:
+    #     import apex
+    #     model, optimizer = apex.amp.initialize(model.cuda(), optimizer, opt_level="O1")
+    #     warnings.warn('**** Initializing mixed precision done. ****')
+
+    # put model on gpus
+    if distributed:
+        find_unused_parameters = cfg.get('find_unused_parameters', False)
+        # Sets the `find_unused_parameters` parameter in
+        # torch.nn.parallel.DistributedDataParallel
+        model = MMDistributedDataParallel(
+            model.cuda(),
+            device_ids=[torch.cuda.current_device()],
+            broadcast_buffers=False,
+            find_unused_parameters=find_unused_parameters)
+    else:
+        model = MMDataParallel(
+            model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)
+
+    if cfg.get('runner') is None:
+        cfg.runner = {
+            'type': 'EpochBasedRunner',
+            'max_epochs': cfg.total_epochs
+        }
+        warnings.warn(
+            'config is now expected to have a `runner` section, '
+            'please set `runner` in your config.', UserWarning)
+
+    runner = build_runner(
+        cfg.runner,
+        default_args=dict(
+            model=model,
+            batch_processor=None,
+            optimizer=optimizer,
+            work_dir=cfg.work_dir,
+            logger=logger,
+            meta=meta))
+
+    # an ugly walkaround to make the .log and .log.json filenames the same
+    runner.timestamp = timestamp
+
+    # fp16 setting
+    fp16_cfg = cfg.get('fp16', None)
+    if fp16_cfg is not None:
+        optimizer_config = Fp16OptimizerHook(
+            **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
+    elif distributed and 'type' not in cfg.optimizer_config:
+        # optimizer_config = DistOptimizerHook(**cfg.optimizer_config)
+        optimizer_config = OptimizerHook(**cfg.optimizer_config)
+    else:
+        optimizer_config = cfg.optimizer_config
+
+    # register hooks
+    runner.register_training_hooks(cfg.lr_config, optimizer_config,
+                                   cfg.checkpoint_config, cfg.log_config,
+                                   cfg.get('momentum_config', None))
+    if distributed and cfg.runner.type=='EpochBasedRunner':
+        runner.register_hook(DistSamplerSeedHook())
+
+    # register eval hooks
+    if validate:
+        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
+        val_dataloader = build_dataloader(
+            val_dataset,
+            samples_per_gpu=cfg.data.samples_per_gpu,
+            workers_per_gpu=cfg.data.workers_per_gpu,
+            dist=distributed,
+            shuffle=False,
+            round_up=True)
+        eval_cfg = cfg.get('evaluation', {})
+        eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
+        eval_hook = DistEvalHook if distributed else EvalHook
+        runner.register_hook(eval_hook(val_dataloader, **eval_cfg), priority='LOW')
+
+    if cfg.resume_from:
+        runner.resume(cfg.resume_from)
+    elif cfg.load_from:
+        runner.load_checkpoint(cfg.load_from)
+    runner.run(data_loaders, cfg.workflow)

CAGE_expression_inference-apvit/apvit_mmcls/core/__init__.py

@@ -0,0 +1,3 @@
+from .evaluation import *  # noqa: F401, F403
+from .fp16 import *  # noqa: F401, F403
+from .utils import *  # noqa: F401, F403

CAGE_expression_inference-apvit/apvit_mmcls/core/evaluation/__init__.py

@@ -0,0 +1,8 @@
+from .eval_hooks import DistEvalHook, EvalHook
+from .mean_ap import average_precision, mAP
+from .multilabel_eval_metrics import average_performance
+
+__all__ = [
+    'DistEvalHook', 'EvalHook', 'average_precision', 'mAP',
+    'average_performance'
+]

CAGE_expression_inference-apvit/apvit_mmcls/core/evaluation/eval_hooks.py

@@ -0,0 +1,158 @@
+import os.path as osp
+from math import inf
+
+from mmcv.runner import Hook
+from torch.utils.data import DataLoader
+from mmcls.utils import get_root_logger
+
+
+class EvalHook(Hook):
+    """Evaluation hook.
+
+    Args:
+        dataloader (DataLoader): A PyTorch dataloader.
+        interval (int): Evaluation interval (by epochs). Default: 1.
+    """
+
+    def __init__(self, dataloader, interval=1, by_epoch=True, **eval_kwargs):
+        if not isinstance(dataloader, DataLoader):
+            raise TypeError('dataloader must be a pytorch DataLoader, but got'
+                            f' {type(dataloader)}')
+        self.dataloader = dataloader
+        self.interval = interval
+        self.eval_kwargs = eval_kwargs
+        self.by_epoch = by_epoch
+        self.logger = get_root_logger()
+
+    def after_train_epoch(self, runner):
+        if not self.by_epoch or not self.every_n_epochs(runner, self.interval):
+            return
+        from mmcls.apis import single_gpu_test
+        results = single_gpu_test(runner.model, self.dataloader, show=False)
+        self.evaluate(runner, results)
+
+    def after_train_iter(self, runner):
+        if self.by_epoch or not self.every_n_iters(runner, self.interval):
+            return
+        from mmcls.apis import single_gpu_test
+        runner.log_buffer.clear()
+        results = single_gpu_test(runner.model, self.dataloader, show=False)
+        self.evaluate(runner, results)
+
+    def evaluate(self, runner, results):
+        eval_res = self.dataloader.dataset.evaluate(
+            results, logger=runner.logger, **self.eval_kwargs)
+        for name, val in eval_res.items():
+            runner.log_buffer.output[name] = val
+        runner.log_buffer.ready = True
+        return eval_res
+
+
+class DistEvalHook(EvalHook):
+    """Distributed evaluation hook.
+
+    Args:
+        dataloader (DataLoader): A PyTorch dataloader.
+        interval (int): Evaluation interval (by epochs). Default: 1.
+        tmpdir (str | None): Temporary directory to save the results of all
+            processes. Default: None.
+        gpu_collect (bool): Whether to use gpu or cpu to collect results.
+            Default: False.
+    """
+
+    def __init__(self,
+                 dataloader,
+                 interval=1,
+                 gpu_collect=True,
+                 by_epoch=True,
+                 print_best=True,
+                 **eval_kwargs):
+        if not isinstance(dataloader, DataLoader):
+            raise TypeError('dataloader must be a pytorch DataLoader, but got '
+                            f'{type(dataloader)}')
+        super().__init__(dataloader, interval, by_epoch, **eval_kwargs)
+        self.dataloader = dataloader
+        self.interval = interval
+        self.gpu_collect = gpu_collect
+        self.by_epoch = by_epoch
+        self.eval_kwargs = eval_kwargs
+        self.print_best = print_best
+    
+    def before_run(self, runner):
+        if self.print_best is not None:
+            if runner.meta is None:
+                warnings.warn('runner.meta is None. Creating a empty one.')
+                runner.meta = dict()
+            runner.meta.setdefault('hook_msgs', dict())
+    
+    def before_train_epoch(self, runner):
+        return
+        # freeze IRNet
+        # frozen_blocks = 7
+        model = runner.model.module.convert
+        model.eval()
+        for param in model.parameters():
+            param.requires_grad = False
+        return
+
+        if frozen_blocks > 0:
+            print(f'IRSE freeze the first {frozen_blocks} blocks, it has {len(model.body)} blocks ')
+            model.input_layer.eval()
+            print('in freeze', model.input_layer[1].training)
+            for param in model.input_layer.parameters():
+                param.requires_grad = False
+        
+        for i in range(frozen_blocks):
+            m = model.body[i]
+            m.eval()
+            for param in m.parameters():
+                param.requires_grad = False
+
+    def after_train_epoch(self, runner):
+        if not self.by_epoch or not self.every_n_epochs(runner, self.interval):
+            return
+        from mmcls.apis import multi_gpu_test
+        results = multi_gpu_test(
+            runner.model,
+            self.dataloader,
+            tmpdir=osp.join(runner.work_dir, '.eval_hook'),
+            gpu_collect=self.gpu_collect)
+        if runner.rank == 0:
+            print('\n')
+            eval_res = self.evaluate(runner, results)
+            if self.print_best:
+                best_score = runner.meta['hook_msgs'].get('best_score', -inf)
+                if 'top-1' not in eval_res:
+                    return
+                acc = eval_res['top-1']
+                if acc > best_score:
+                    self.logger.info(f'top-1 accuracy improved from {best_score} to {acc}')
+                    runner.meta['hook_msgs']['best_score'] = acc
+                    runner.save_checkpoint(runner.work_dir, save_optimizer=False, filename_tmpl='best.pth', create_symlink=False)
+                else:
+                    self.logger.info(f'top-1 accuracy did not improve from {best_score}')
+
+    def after_train_iter(self, runner):
+        if self.by_epoch or not self.every_n_iters(runner, self.interval):
+            return
+        from mmcls.apis import multi_gpu_test
+        runner.log_buffer.clear()
+        results = multi_gpu_test(
+            runner.model,
+            self.dataloader,
+            tmpdir=osp.join(runner.work_dir, '.eval_hook'),
+            gpu_collect=self.gpu_collect)
+        if runner.rank == 0:
+            print('\n')
+            eval_res = self.evaluate(runner, results)
+            if self.print_best:
+                best_score = runner.meta['hook_msgs'].get('best_score', -inf)
+                if 'top-1' not in eval_res:
+                    return
+                acc = eval_res['top-1']
+                if acc > best_score:
+                    self.logger.info(f'top-1 accuracy improved from {best_score} to {acc}')
+                    runner.meta['hook_msgs']['best_score'] = acc
+                    runner.save_checkpoint(runner.work_dir, save_optimizer=False, filename_tmpl='best.pth', create_symlink=False)
+                else:
+                    self.logger.info(f'top-1 accuracy did not improve from {best_score}')

CAGE_expression_inference-apvit/apvit_mmcls/core/evaluation/mean_ap.py

@@ -0,0 +1,73 @@
+import numpy as np
+import torch
+
+
+def average_precision(pred, target):
+    """ Calculate the average precision for a single class
+
+    AP summarizes a precision-recall curve as the weighted mean of maximum
+    precisions obtained for any r'>r, where r is the recall:
+
+    ..math::
+        \\text{AP} = \\sum_n (R_n - R_{n-1}) P_n
+
+    Note that no approximation is involved since the curve is piecewise
+    constant.
+
+    Args:
+        pred (np.ndarray): The model prediction with shape (N, ).
+        target (np.ndarray): The target of each prediction with shape (N, ).
+
+    Returns:
+        float: a single float as average precision value.
+    """
+    eps = np.finfo(np.float32).eps
+
+    # sort examples
+    sort_inds = np.argsort(-pred)
+    sort_target = target[sort_inds]
+
+    # count true positive examples
+    pos_inds = sort_target == 1
+    tp = np.cumsum(pos_inds)
+    total_pos = tp[-1]
+
+    # count not difficult examples
+    pn_inds = sort_target != -1
+    pn = np.cumsum(pn_inds)
+
+    tp[np.logical_not(pos_inds)] = 0
+    precision = tp / np.maximum(pn, eps)
+    ap = np.sum(precision) / np.maximum(total_pos, eps)
+    return ap
+
+
+def mAP(pred, target):
+    """ Calculate the mean average precision with respect of classes
+
+    Args:
+        pred (torch.Tensor | np.ndarray): The model prediction with shape
+            (N, C), where C is the number of classes.
+        target (torch.Tensor | np.ndarray): The target of each prediction with
+            shape (N, C), where C is the number of classes. 1 stands for
+            positive examples, 0 stands for negative examples and -1 stands for
+            difficult examples.
+
+    Returns:
+        float: A single float as mAP value.
+    """
+    if isinstance(pred, torch.Tensor) and isinstance(target, torch.Tensor):
+        pred = pred.numpy()
+        target = target.numpy()
+    elif not (isinstance(pred, np.ndarray) and isinstance(target, np.ndarray)):
+        raise TypeError('pred and target should both be torch.Tensor or'
+                        'np.ndarray')
+
+    assert pred.shape == \
+        target.shape, 'pred and target should be in the same shape.'
+    num_classes = pred.shape[1]
+    ap = np.zeros(num_classes)
+    for k in range(num_classes):
+        ap[k] = average_precision(pred[:, k], target[:, k])
+    mean_ap = ap.mean() * 100.0
+    return mean_ap

CAGE_expression_inference-apvit/apvit_mmcls/core/evaluation/multilabel_eval_metrics.py

@@ -0,0 +1,71 @@
+import warnings
+
+import numpy as np
+import torch
+
+
+def average_performance(pred, target, thr=None, k=None):
+    """Calculate CP, CR, CF1, OP, OR, OF1, where C stands for per-class
+        average, O stands for overall average, P stands for precision, R
+        stands for recall and F1 stands for F1-score
+
+    Args:
+        pred (torch.Tensor | np.ndarray): The model prediction with shape
+            (N, C), where C is the number of classes.
+        target (torch.Tensor | np.ndarray): The target of each prediction with
+            shape (N, C), where C is the number of classes. 1 stands for
+            positive examples, 0 stands for negative examples and -1 stands for
+            difficult examples.
+        thr (float): The confidence threshold. Defaults to None.
+        k (int): Top-k performance. Note that if thr and k are both given, k
+            will be ignored. Defaults to None.
+
+    Returns:
+        tuple: (CP, CR, CF1, OP, OR, OF1)
+    """
+    if isinstance(pred, torch.Tensor) and isinstance(target, torch.Tensor):
+        pred = pred.numpy()
+        target = target.numpy()
+    elif not (isinstance(pred, np.ndarray) and isinstance(target, np.ndarray)):
+        raise TypeError('pred and target should both be torch.Tensor or'
+                        'np.ndarray')
+    if thr is None and k is None:
+        thr = 0.5
+        warnings.warn('Neither thr nor k is given, set thr as 0.5 by '
+                      'default.')
+    elif thr is not None and k is not None:
+        warnings.warn('Both thr and k are given, use threshold in favor of '
+                      'top-k.')
+
+    assert pred.shape == \
+        target.shape, 'pred and target should be in the same shape.'
+
+    eps = np.finfo(np.float32).eps
+    target[target == -1] = 0
+    if thr is not None:
+        # a label is predicted positive if the confidence is no lower than thr
+        pos_inds = pred >= thr
+
+    else:
+        # top-k labels will be predicted positive for any example
+        sort_inds = np.argsort(-pred, axis=1)
+        sort_inds_ = sort_inds[:, :k]
+        inds = np.indices(sort_inds_.shape)
+        pos_inds = np.zeros_like(pred)
+        pos_inds[inds[0], sort_inds_] = 1
+
+    tp = (pos_inds * target) == 1
+    fp = (pos_inds * (1 - target)) == 1
+    fn = ((1 - pos_inds) * target) == 1
+
+    precision_class = tp.sum(axis=0) / np.maximum(
+        tp.sum(axis=0) + fp.sum(axis=0), eps)
+    recall_class = tp.sum(axis=0) / np.maximum(
+        tp.sum(axis=0) + fn.sum(axis=0), eps)
+    CP = precision_class.mean() * 100.0
+    CR = recall_class.mean() * 100.0
+    CF1 = 2 * CP * CR / np.maximum(CP + CR, eps)
+    OP = tp.sum() / np.maximum(tp.sum() + fp.sum(), eps) * 100.0
+    OR = tp.sum() / np.maximum(tp.sum() + fn.sum(), eps) * 100.0
+    OF1 = 2 * OP * OR / np.maximum(OP + OR, eps)
+    return CP, CR, CF1, OP, OR, OF1

CAGE_expression_inference-apvit/apvit_mmcls/core/fp16/__init__.py

@@ -0,0 +1,4 @@
+from .decorators import auto_fp16, force_fp32
+from .hooks import Fp16OptimizerHook, wrap_fp16_model
+
+__all__ = ['auto_fp16', 'force_fp32', 'Fp16OptimizerHook', 'wrap_fp16_model']

CAGE_expression_inference-apvit/apvit_mmcls/core/fp16/decorators.py

@@ -0,0 +1,160 @@
+import functools
+from inspect import getfullargspec
+
+import torch
+
+from .utils import cast_tensor_type
+
+
+def auto_fp16(apply_to=None, out_fp32=False):
+    """Decorator to enable fp16 training automatically.
+
+    This decorator is useful when you write custom modules and want to support
+    mixed precision training. If inputs arguments are fp32 tensors, they will
+    be converted to fp16 automatically. Arguments other than fp32 tensors are
+    ignored.
+
+    Args:
+        apply_to (Iterable, optional): The argument names to be converted.
+            `None` indicates all arguments.
+        out_fp32 (bool): Whether to convert the output back to fp32.
+
+    :Example:
+
+        class MyModule1(nn.Module)
+
+            # Convert x and y to fp16
+            @auto_fp16()
+            def forward(self, x, y):
+                pass
+
+        class MyModule2(nn.Module):
+
+            # convert pred to fp16
+            @auto_fp16(apply_to=('pred', ))
+            def do_something(self, pred, others):
+                pass
+    """
+
+    def auto_fp16_wrapper(old_func):
+
+        @functools.wraps(old_func)
+        def new_func(*args, **kwargs):
+            # check if the module has set the attribute `fp16_enabled`, if not,
+            # just fallback to the original method.
+            if not isinstance(args[0], torch.nn.Module):
+                raise TypeError('@auto_fp16 can only be used to decorate the '
+                                'method of nn.Module')
+            if not (hasattr(args[0], 'fp16_enabled') and args[0].fp16_enabled):
+                return old_func(*args, **kwargs)
+            # get the arg spec of the decorated method
+            args_info = getfullargspec(old_func)
+            # get the argument names to be casted
+            args_to_cast = args_info.args if apply_to is None else apply_to
+            # convert the args that need to be processed
+            new_args = []
+            # NOTE: default args are not taken into consideration
+            if args:
+                arg_names = args_info.args[:len(args)]
+                for i, arg_name in enumerate(arg_names):
+                    if arg_name in args_to_cast:
+                        new_args.append(
+                            cast_tensor_type(args[i], torch.float, torch.half))
+                    else:
+                        new_args.append(args[i])
+            # convert the kwargs that need to be processed
+            new_kwargs = {}
+            if kwargs:
+                for arg_name, arg_value in kwargs.items():
+                    if arg_name in args_to_cast:
+                        new_kwargs[arg_name] = cast_tensor_type(
+                            arg_value, torch.float, torch.half)
+                    else:
+                        new_kwargs[arg_name] = arg_value
+            # apply converted arguments to the decorated method
+            output = old_func(*new_args, **new_kwargs)
+            # cast the results back to fp32 if necessary
+            if out_fp32:
+                output = cast_tensor_type(output, torch.half, torch.float)
+            return output
+
+        return new_func
+
+    return auto_fp16_wrapper
+
+
+def force_fp32(apply_to=None, out_fp16=False):
+    """Decorator to convert input arguments to fp32 in force.
+
+    This decorator is useful when you write custom modules and want to support
+    mixed precision training. If there are some inputs that must be processed
+    in fp32 mode, then this decorator can handle it. If inputs arguments are
+    fp16 tensors, they will be converted to fp32 automatically. Arguments other
+    than fp16 tensors are ignored.
+
+    Args:
+        apply_to (Iterable, optional): The argument names to be converted.
+            `None` indicates all arguments.
+        out_fp16 (bool): Whether to convert the output back to fp16.
+
+    :Example:
+
+        class MyModule1(nn.Module)
+
+            # Convert x and y to fp32
+            @force_fp32()
+            def loss(self, x, y):
+                pass
+
+        class MyModule2(nn.Module):
+
+            # convert pred to fp32
+            @force_fp32(apply_to=('pred', ))
+            def post_process(self, pred, others):
+                pass
+    """
+
+    def force_fp32_wrapper(old_func):
+
+        @functools.wraps(old_func)
+        def new_func(*args, **kwargs):
+            # check if the module has set the attribute `fp16_enabled`, if not,
+            # just fallback to the original method.
+            if not isinstance(args[0], torch.nn.Module):
+                raise TypeError('@force_fp32 can only be used to decorate the '
+                                'method of nn.Module')
+            if not (hasattr(args[0], 'fp16_enabled') and args[0].fp16_enabled):
+                return old_func(*args, **kwargs)
+            # get the arg spec of the decorated method
+            args_info = getfullargspec(old_func)
+            # get the argument names to be casted
+            args_to_cast = args_info.args if apply_to is None else apply_to
+            # convert the args that need to be processed
+            new_args = []
+            if args:
+                arg_names = args_info.args[:len(args)]
+                for i, arg_name in enumerate(arg_names):
+                    if arg_name in args_to_cast:
+                        new_args.append(
+                            cast_tensor_type(args[i], torch.half, torch.float))
+                    else:
+                        new_args.append(args[i])
+            # convert the kwargs that need to be processed
+            new_kwargs = dict()
+            if kwargs:
+                for arg_name, arg_value in kwargs.items():
+                    if arg_name in args_to_cast:
+                        new_kwargs[arg_name] = cast_tensor_type(
+                            arg_value, torch.half, torch.float)
+                    else:
+                        new_kwargs[arg_name] = arg_value
+            # apply converted arguments to the decorated method
+            output = old_func(*new_args, **new_kwargs)
+            # cast the results back to fp32 if necessary
+            if out_fp16:
+                output = cast_tensor_type(output, torch.float, torch.half)
+            return output
+
+        return new_func
+
+    return force_fp32_wrapper

CAGE_expression_inference-apvit/apvit_mmcls/core/fp16/hooks.py

@@ -0,0 +1,132 @@
+import copy
+
+import torch
+import torch.nn as nn
+from mmcv.runner import OptimizerHook
+from mmcv.utils.parrots_wrapper import _BatchNorm
+
+from ..utils import allreduce_grads
+from .utils import cast_tensor_type
+
+
+class Fp16OptimizerHook(OptimizerHook):
+    """FP16 optimizer hook.
+
+    The steps of fp16 optimizer is as follows.
+    1. Scale the loss value.
+    2. BP in the fp16 model.
+    2. Copy gradients from fp16 model to fp32 weights.
+    3. Update fp32 weights.
+    4. Copy updated parameters from fp32 weights to fp16 model.
+
+    Refer to https://arxiv.org/abs/1710.03740 for more details.
+
+    Args:
+        loss_scale (float): Scale factor multiplied with loss.
+    """
+
+    def __init__(self,
+                 grad_clip=None,
+                 coalesce=True,
+                 bucket_size_mb=-1,
+                 loss_scale=512.,
+                 distributed=True):
+        self.grad_clip = grad_clip
+        self.coalesce = coalesce
+        self.bucket_size_mb = bucket_size_mb
+        self.loss_scale = loss_scale
+        self.distributed = distributed
+
+    def before_run(self, runner):
+        # keep a copy of fp32 weights
+        runner.optimizer.param_groups = copy.deepcopy(
+            runner.optimizer.param_groups)
+        # convert model to fp16
+        wrap_fp16_model(runner.model)
+
+    def copy_grads_to_fp32(self, fp16_net, fp32_weights):
+        """Copy gradients from fp16 model to fp32 weight copy."""
+        for fp32_param, fp16_param in zip(fp32_weights, fp16_net.parameters()):
+            if fp16_param.grad is not None:
+                if fp32_param.grad is None:
+                    fp32_param.grad = fp32_param.data.new(fp32_param.size())
+                fp32_param.grad.copy_(fp16_param.grad)
+
+    def copy_params_to_fp16(self, fp16_net, fp32_weights):
+        """Copy updated params from fp32 weight copy to fp16 model."""
+        for fp16_param, fp32_param in zip(fp16_net.parameters(), fp32_weights):
+            fp16_param.data.copy_(fp32_param.data)
+
+    def after_train_iter(self, runner):
+        # clear grads of last iteration
+        runner.model.zero_grad()
+        runner.optimizer.zero_grad()
+        # scale the loss value
+        scaled_loss = runner.outputs['loss'] * self.loss_scale
+        scaled_loss.backward()
+        # copy fp16 grads in the model to fp32 params in the optimizer
+        fp32_weights = []
+        for param_group in runner.optimizer.param_groups:
+            fp32_weights += param_group['params']
+        self.copy_grads_to_fp32(runner.model, fp32_weights)
+        # allreduce grads
+        if self.distributed:
+            allreduce_grads(fp32_weights, self.coalesce, self.bucket_size_mb)
+        # scale the gradients back
+        for param in fp32_weights:
+            if param.grad is not None:
+                param.grad.div_(self.loss_scale)
+        if self.grad_clip is not None:
+            grad_norm = self.clip_grads(fp32_weights)
+            if grad_norm is not None:
+                # Add grad norm to the logger
+                runner.log_buffer.update({'grad_norm': float(grad_norm)},
+                                            runner.outputs['num_samples'])
+        # update fp32 params
+        runner.optimizer.step()
+        # copy fp32 params to the fp16 model
+        self.copy_params_to_fp16(runner.model, fp32_weights)
+
+
+def wrap_fp16_model(model):
+    # convert model to fp16
+    model.half()
+    # patch the normalization layers to make it work in fp32 mode
+    patch_norm_fp32(model)
+    # set `fp16_enabled` flag
+    for m in model.modules():
+        if hasattr(m, 'fp16_enabled'):
+            m.fp16_enabled = True
+
+
+def patch_norm_fp32(module):
+    if isinstance(module, (_BatchNorm, nn.GroupNorm)):
+        module.float()
+        module.forward = patch_forward_method(module.forward, torch.half,
+                                              torch.float)
+    for child in module.children():
+        patch_norm_fp32(child)
+    return module
+
+
+def patch_forward_method(func, src_type, dst_type, convert_output=True):
+    """Patch the forward method of a module.
+
+    Args:
+        func (callable): The original forward method.
+        src_type (torch.dtype): Type of input arguments to be converted from.
+        dst_type (torch.dtype): Type of input arguments to be converted to.
+        convert_output (bool): Whether to convert the output back to src_type.
+
+    Returns:
+        callable: The patched forward method.
+    """
+
+    def new_forward(*args, **kwargs):
+        output = func(*cast_tensor_type(args, src_type, dst_type),
+                      **cast_tensor_type(kwargs, src_type, dst_type))
+        if convert_output:
+            output = cast_tensor_type(output, dst_type, src_type)
+        return output
+
+    return new_forward

CAGE_expression_inference-apvit/apvit_mmcls/core/fp16/utils.py

@@ -0,0 +1,23 @@
+from collections import abc
+
+import numpy as np
+import torch
+
+
+def cast_tensor_type(inputs, src_type, dst_type):
+    if isinstance(inputs, torch.Tensor):
+        return inputs.to(dst_type)
+    elif isinstance(inputs, str):
+        return inputs
+    elif isinstance(inputs, np.ndarray):
+        return inputs
+    elif isinstance(inputs, abc.Mapping):
+        return type(inputs)({
+            k: cast_tensor_type(v, src_type, dst_type)
+            for k, v in inputs.items()
+        })
+    elif isinstance(inputs, abc.Iterable):
+        return type(inputs)(
+            cast_tensor_type(item, src_type, dst_type) for item in inputs)
+    else:
+        return inputs

CAGE_expression_inference-apvit/apvit_mmcls/core/utils/__init__.py

@@ -0,0 +1,4 @@
+from .dist_utils import DistOptimizerHook, allreduce_grads
+from .misc import multi_apply
+
+__all__ = ['allreduce_grads', 'DistOptimizerHook', 'multi_apply']

CAGE_expression_inference-apvit/apvit_mmcls/core/utils/dist_utils.py

@@ -0,0 +1,56 @@
+from collections import OrderedDict
+
+import torch.distributed as dist
+from mmcv.runner import OptimizerHook
+from torch._utils import (_flatten_dense_tensors, _take_tensors,
+                          _unflatten_dense_tensors)
+
+
+def _allreduce_coalesced(tensors, world_size, bucket_size_mb=-1):
+    if bucket_size_mb > 0:
+        bucket_size_bytes = bucket_size_mb * 1024 * 1024
+        buckets = _take_tensors(tensors, bucket_size_bytes)
+    else:
+        buckets = OrderedDict()
+        for tensor in tensors:
+            tp = tensor.type()
+            if tp not in buckets:
+                buckets[tp] = []
+            buckets[tp].append(tensor)
+        buckets = buckets.values()
+
+    for bucket in buckets:
+        flat_tensors = _flatten_dense_tensors(bucket)
+        dist.all_reduce(flat_tensors)
+        flat_tensors.div_(world_size)
+        for tensor, synced in zip(
+                bucket, _unflatten_dense_tensors(flat_tensors, bucket)):
+            tensor.copy_(synced)
+
+
+def allreduce_grads(params, coalesce=True, bucket_size_mb=-1):
+    grads = [
+        param.grad.data for param in params
+        if param.requires_grad and param.grad is not None
+    ]
+    world_size = dist.get_world_size()
+    if coalesce:
+        _allreduce_coalesced(grads, world_size, bucket_size_mb)
+    else:
+        for tensor in grads:
+            dist.all_reduce(tensor.div_(world_size))
+
+
+class DistOptimizerHook(OptimizerHook):
+
+    def __init__(self, grad_clip=None, coalesce=True, bucket_size_mb=-1):
+        self.grad_clip = grad_clip
+        self.coalesce = coalesce
+        self.bucket_size_mb = bucket_size_mb
+
+    def after_train_iter(self, runner):
+        runner.optimizer.zero_grad()
+        runner.outputs['loss'].backward()
+        if self.grad_clip is not None:
+            self.clip_grads(runner.model.parameters())
+        runner.optimizer.step()

CAGE_expression_inference-apvit/apvit_mmcls/core/utils/misc.py

@@ -0,0 +1,7 @@
+from functools import partial
+
+
+def multi_apply(func, *args, **kwargs):
+    pfunc = partial(func, **kwargs) if kwargs else func
+    map_results = map(pfunc, *args)
+    return tuple(map(list, zip(*map_results)))

CAGE_expression_inference-apvit/apvit_mmcls/datasets/__init__.py

@@ -0,0 +1,16 @@
+from .base_dataset import BaseDataset
+from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset
+from .cifar import CIFAR10, CIFAR100
+from .dataset_wrappers import (ClassBalancedDataset, ConcatDataset,
+                               RepeatDataset)
+from .imagenet import ImageNet
+from .mnist import MNIST, FashionMNIST
+from .samplers import DistributedSampler
+from .raf import RAF
+
+__all__ = [
+    'BaseDataset', 'ImageNet', 'CIFAR10', 'CIFAR100', 'MNIST', 'FashionMNIST',
+    'build_dataloader', 'build_dataset', 'Compose', 'DistributedSampler',
+    'ConcatDataset', 'RepeatDataset', 'ClassBalancedDataset', 'DATASETS',
+    'PIPELINES', 'RAF', 
+]

CAGE_expression_inference-apvit/apvit_mmcls/datasets/base_dataset.py

@@ -0,0 +1,185 @@
+import copy
+from abc import ABCMeta, abstractmethod
+import os
+from shutil import copyfile
+
+import mmcv
+import numpy as np
+from torch.utils.data import Dataset
+
+from mmcls.models.losses import accuracy, f1_score, precision, recall
+from mmcls.models.losses.eval_metrics import class_accuracy
+from .pipelines import Compose
+
+
+class BaseDataset(Dataset, metaclass=ABCMeta):
+    """Base dataset.
+
+    Args:
+        data_prefix (str): the prefix of data path
+        pipeline (list): a list of dict, where each element represents
+            a operation defined in `mmcls.datasets.pipelines`
+        ann_file (str | None): the annotation file. When ann_file is str,
+            the subclass is expected to read from the ann_file. When ann_file
+            is None, the subclass is expected to read according to data_prefix
+        test_mode (bool): in train mode or test mode
+    """
+
+    CLASSES = None
+
+    def __init__(self,
+                 data_prefix,
+                 pipeline,
+                 classes=None,
+                 ann_file=None,
+                 test_mode=False):
+        super(BaseDataset, self).__init__()
+
+        self.ann_file = ann_file
+        self.data_prefix = data_prefix
+        self.test_mode = test_mode
+        self.pipeline = Compose(pipeline)
+        self.data_infos = self.load_annotations()
+        # self.data_infos = self.load_exp_and_au_annotations()
+        self.CLASSES = self.get_classes(classes)
+        if os.environ.get('DEBUG_MODE', '0') == '1':
+            self.data_infos = self.data_infos[:30]
+
+    @abstractmethod
+    def load_annotations(self):
+        pass
+
+    @property
+    def class_to_idx(self):
+        """Map mapping class name to class index.
+
+        Returns:
+            dict: mapping from class name to class index.
+        """
+
+        return {_class: i for i, _class in enumerate(self.CLASSES)}
+
+    def get_gt_labels(self):
+        """Get all ground-truth labels (categories).
+
+        Returns:
+            list[int]: categories for all images.
+        """
+
+        gt_labels = np.array([data['gt_label'] for data in self.data_infos])
+        return gt_labels
+    
+    def get_coarse_labels(self):
+        coarse_labels = np.array([data['coarse_label'] for data in self.data_infos])
+        return coarse_labels
+
+    def get_cat_ids(self, idx):
+        """Get category id by index.
+
+        Args:
+            idx (int): Index of data.
+
+        Returns:
+            int: Image category of specified index.
+        """
+
+        return (int(self.data_infos[idx]['gt_label'].astype(np.int)), )
+
+    def prepare_data(self, idx):
+        results = copy.deepcopy(self.data_infos[idx])
+        return self.pipeline(results)
+
+    def __len__(self):
+        return len(self.data_infos)
+
+    def __getitem__(self, idx):
+        return self.prepare_data(idx)
+
+    @classmethod
+    def get_classes(cls, classes=None):
+        """Get class names of current dataset.
+        Args:
+            classes (Sequence[str] | str | None): If classes is None, use
+                default CLASSES defined by builtin dataset. If classes is a
+                string, take it as a file name. The file contains the name of
+                classes where each line contains one class name. If classes is
+                a tuple or list, override the CLASSES defined by the dataset.
+
+        Returns:
+            tuple[str] or list[str]: Names of categories of the dataset.
+        """
+        if classes is None:
+            return cls.CLASSES
+
+        if isinstance(classes, str):
+            # take it as a file path
+            class_names = mmcv.list_from_file(classes)
+        elif isinstance(classes, (tuple, list)):
+            class_names = classes
+        else:
+            raise ValueError(f'Unsupported type {type(classes)} of classes.')
+
+        return class_names
+
+    def evaluate(self,
+                results,
+                metric='accuracy',
+                metric_options={'topk': (1, 2)},
+                logger=None):
+        """Evaluate the dataset.
+
+        Args:
+            results (list): Testing results of the dataset.
+            metric (str | list[str]): Metrics to be evaluated.
+                Default value is `accuracy`.
+            logger (logging.Logger | None | str): Logger used for printing
+                related information during evaluation. Default: None.
+        Returns:
+            dict: evaluation results
+        """
+        if isinstance(metric, str):
+            metrics = [metric]
+        else:
+            metrics = metric
+        allowed_metrics = ['accuracy', 'precision', 'recall', 'f1_score', 'class_accuracy']
+        eval_results = {}
+        for metric in metrics:
+            if metric not in allowed_metrics:
+                raise KeyError(f'metric {metric} is not supported.')
+            results = np.vstack(results)
+            gt_labels = self.get_gt_labels()
+            num_imgs = len(results)
+            assert len(gt_labels) == num_imgs, f'{len(gt_labels)}, {num_imgs}'
+            # convert gt_labels to another dataset format for cross-dataset test
+            # raf2affect_dict = [6, 3, 4, 0, 1, 2, 5] # RAF -> AffectNet实验需要将RAF标签修改为AffectNet的标签
+            # affect2raf_dict = [3, 4, 5, 1, 2, 6, 0]
+            # raf2ferplus_dict = [6,3,0,4,5,2,1,None]
+            # gt_labels = [raf2ferplus_dict[i] for i in gt_labels]
+            # # specific process for RAF -> FERPlus
+            # results = [results[i] for i in range(len(results)) if gt_labels[i] is not None]
+            # gt_labels = [gt_labels[i] for i in range(len(gt_labels)) if gt_labels[i] is not None]
+            # results = np.array(results)
+            # gt_labels = np.array(gt_labels)
+
+            if metric == 'accuracy':
+                topk = metric_options.get('topk')
+                acc = accuracy(results, gt_labels, topk)
+                eval_result = {f'top-{k}': a.item() for k, a in zip(topk, acc)}
+            elif metric == 'precision':
+                precision_value = precision(results, gt_labels)
+                eval_result = {'precision': precision_value}
+            elif metric == 'recall':
+                recall_value = recall(results, gt_labels)
+                eval_result = {'recall': recall_value}
+            elif metric == 'f1_score':
+                f1_score_value = f1_score(results, gt_labels)
+                eval_result = {'f1_score': f1_score_value}
+            elif metric == 'class_accuracy':
+                class_accuracy_value = class_accuracy(results, gt_labels, self.CLASSES)
+                print('\n')
+                for name, val in zip(self.CLASSES, class_accuracy_value):
+                    print(f'{name}: \t{val}')
+                print('\n')
+                eval_result = dict()
+            eval_results.update(eval_result)
+        return eval_results

CAGE_expression_inference-apvit/apvit_mmcls/datasets/builder.py

@@ -0,0 +1,110 @@
+import platform
+import random
+from functools import partial
+
+import numpy as np
+from mmcv.parallel import collate
+from mmcv.runner import get_dist_info
+from mmcv.utils import Registry, build_from_cfg
+from torch.utils.data import DataLoader
+
+from .samplers import DistributedSampler
+
+if platform.system() != 'Windows':
+    # https://github.com/pytorch/pytorch/issues/973
+    import resource
+    rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
+    hard_limit = rlimit[1]
+    soft_limit = min(4096, hard_limit)
+    resource.setrlimit(resource.RLIMIT_NOFILE, (soft_limit, hard_limit))
+
+DATASETS = Registry('dataset')
+PIPELINES = Registry('pipeline')
+
+
+def build_dataset(cfg, default_args=None):
+    from .dataset_wrappers import (ConcatDataset, RepeatDataset,
+                                   ClassBalancedDataset)
+    cfg = cfg.copy()
+    if isinstance(cfg, (list, tuple)):
+        dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg])
+    elif cfg['type'] == 'RepeatDataset':
+        dataset = RepeatDataset(
+            build_dataset(cfg['dataset'], default_args), cfg['times'])
+    elif cfg['type'] == 'ClassBalancedDataset':
+        ds = build_dataset(cfg['dataset'], default_args)
+        cfg['dataset'] = ds
+        dataset = ClassBalancedDataset(**cfg)
+    else:
+        dataset = build_from_cfg(cfg, DATASETS, default_args)
+
+    return dataset
+
+
+def build_dataloader(dataset,
+                     samples_per_gpu,
+                     workers_per_gpu,
+                     num_gpus=1,
+                     dist=True,
+                     shuffle=True,
+                     round_up=True,
+                     seed=None,
+                     **kwargs):
+    """Build PyTorch DataLoader.
+
+    In distributed training, each GPU/process has a dataloader.
+    In non-distributed training, there is only one dataloader for all GPUs.
+
+    Args:
+        dataset (Dataset): A PyTorch dataset.
+        samples_per_gpu (int): Number of training samples on each GPU, i.e.,
+            batch size of each GPU.
+        workers_per_gpu (int): How many subprocesses to use for data loading
+            for each GPU.
+        num_gpus (int): Number of GPUs. Only used in non-distributed training.
+        dist (bool): Distributed training/test or not. Default: True.
+        shuffle (bool): Whether to shuffle the data at every epoch.
+            Default: True.
+        round_up (bool): Whether to round up the length of dataset by adding
+            extra samples to make it evenly divisible. Default: True.
+        kwargs: any keyword argument to be used to initialize DataLoader
+
+    Returns:
+        DataLoader: A PyTorch dataloader.
+    """
+    rank, world_size = get_dist_info()
+    if dist:
+        sampler = DistributedSampler(
+            dataset, world_size, rank, shuffle=shuffle, round_up=round_up)
+        shuffle = False
+        batch_size = samples_per_gpu
+        num_workers = workers_per_gpu
+    else:
+        sampler = None
+        batch_size = num_gpus * samples_per_gpu
+        num_workers = num_gpus * workers_per_gpu
+
+    init_fn = partial(
+        worker_init_fn, num_workers=num_workers, rank=rank,
+        seed=seed) if seed is not None else None
+
+    data_loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        sampler=sampler,
+        num_workers=num_workers,
+        collate_fn=partial(collate, samples_per_gpu=samples_per_gpu),
+        pin_memory=False,
+        shuffle=shuffle,
+        worker_init_fn=init_fn,
+        **kwargs)
+
+    return data_loader
+
+
+def worker_init_fn(worker_id, num_workers, rank, seed):
+    # The seed of each worker equals to
+    # num_worker * rank + worker_id + user_seed
+    worker_seed = num_workers * rank + worker_id + seed
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)

CAGE_expression_inference-apvit/apvit_mmcls/datasets/cifar.py

@@ -0,0 +1,123 @@
+import os
+import os.path
+import pickle
+
+import numpy as np
+
+from .base_dataset import BaseDataset
+from .builder import DATASETS
+from .utils import check_integrity, download_and_extract_archive
+
+
+@DATASETS.register_module()
+class CIFAR10(BaseDataset):
+    """`CIFAR10 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.
+
+    This implementation is modified from
+    https://github.com/pytorch/vision/blob/master/torchvision/datasets/cifar.py  # noqa: E501
+    """
+
+    base_folder = 'cifar-10-batches-py'
+    url = 'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'
+    filename = 'cifar-10-python.tar.gz'
+    tgz_md5 = 'c58f30108f718f92721af3b95e74349a'
+    train_list = [
+        ['data_batch_1', 'c99cafc152244af753f735de768cd75f'],
+        ['data_batch_2', 'd4bba439e000b95fd0a9bffe97cbabec'],
+        ['data_batch_3', '54ebc095f3ab1f0389bbae665268c751'],
+        ['data_batch_4', '634d18415352ddfa80567beed471001a'],
+        ['data_batch_5', '482c414d41f54cd18b22e5b47cb7c3cb'],
+    ]
+
+    test_list = [
+        ['test_batch', '40351d587109b95175f43aff81a1287e'],
+    ]
+    meta = {
+        'filename': 'batches.meta',
+        'key': 'label_names',
+        'md5': '5ff9c542aee3614f3951f8cda6e48888',
+    }
+
+    def load_annotations(self):
+
+        if not self._check_integrity():
+            download_and_extract_archive(
+                self.url,
+                self.data_prefix,
+                filename=self.filename,
+                md5=self.tgz_md5)
+
+        if not self.test_mode:
+            downloaded_list = self.train_list
+        else:
+            downloaded_list = self.test_list
+
+        self.imgs = []
+        self.gt_labels = []
+
+        # load the picked numpy arrays
+        for file_name, checksum in downloaded_list:
+            file_path = os.path.join(self.data_prefix, self.base_folder,
+                                     file_name)
+            with open(file_path, 'rb') as f:
+                entry = pickle.load(f, encoding='latin1')
+                self.imgs.append(entry['data'])
+                if 'labels' in entry:
+                    self.gt_labels.extend(entry['labels'])
+                else:
+                    self.gt_labels.extend(entry['fine_labels'])
+
+        self.imgs = np.vstack(self.imgs).reshape(-1, 3, 32, 32)
+        self.imgs = self.imgs.transpose((0, 2, 3, 1))  # convert to HWC
+
+        self._load_meta()
+
+        data_infos = []
+        for img, gt_label in zip(self.imgs, self.gt_labels):
+            gt_label = np.array(gt_label, dtype=np.int64)
+            info = {'img': img, 'gt_label': gt_label}
+            data_infos.append(info)
+        return data_infos
+
+    def _load_meta(self):
+        path = os.path.join(self.data_prefix, self.base_folder,
+                            self.meta['filename'])
+        if not check_integrity(path, self.meta['md5']):
+            raise RuntimeError(
+                'Dataset metadata file not found or corrupted.' +
+                ' You can use download=True to download it')
+        with open(path, 'rb') as infile:
+            data = pickle.load(infile, encoding='latin1')
+            self.CLASSES = data[self.meta['key']]
+
+    def _check_integrity(self):
+        root = self.data_prefix
+        for fentry in (self.train_list + self.test_list):
+            filename, md5 = fentry[0], fentry[1]
+            fpath = os.path.join(root, self.base_folder, filename)
+            if not check_integrity(fpath, md5):
+                return False
+        return True
+
+
+@DATASETS.register_module()
+class CIFAR100(CIFAR10):
+    """`CIFAR100 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.
+    """
+
+    base_folder = 'cifar-100-python'
+    url = 'https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz'
+    filename = 'cifar-100-python.tar.gz'
+    tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
+    train_list = [
+        ['train', '16019d7e3df5f24257cddd939b257f8d'],
+    ]
+
+    test_list = [
+        ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'],
+    ]
+    meta = {
+        'filename': 'meta',
+        'key': 'fine_label_names',
+        'md5': '7973b15100ade9c7d40fb424638fde48',
+    }

CAGE_expression_inference-apvit/apvit_mmcls/datasets/dataset_wrappers.py

@@ -0,0 +1,173 @@
+import bisect
+import math
+from collections import defaultdict
+
+import numpy as np
+from torch.utils.data.dataset import ConcatDataset as _ConcatDataset
+
+from .builder import DATASETS
+
+
+@DATASETS.register_module()
+class ConcatDataset(_ConcatDataset):
+    """A wrapper of concatenated dataset.
+
+    Same as :obj:`torch.utils.data.dataset.ConcatDataset`, but
+    add `get_cat_ids` function.
+
+    Args:
+        datasets (list[:obj:`Dataset`]): A list of datasets.
+    """
+
+    def __init__(self, datasets):
+        super(ConcatDataset, self).__init__(datasets)
+        self.CLASSES = datasets[0].CLASSES
+
+    def get_cat_ids(self, idx):
+        if idx < 0:
+            if -idx > len(self):
+                raise ValueError(
+                    'absolute value of index should not exceed dataset length')
+            idx = len(self) + idx
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+        if dataset_idx == 0:
+            sample_idx = idx
+        else:
+            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+        return self.datasets[dataset_idx].get_cat_ids(sample_idx)
+
+
+@DATASETS.register_module()
+class RepeatDataset(object):
+    """A wrapper of repeated dataset.
+
+    The length of repeated dataset will be `times` larger than the original
+    dataset. This is useful when the data loading time is long but the dataset
+    is small. Using RepeatDataset can reduce the data loading time between
+    epochs.
+
+    Args:
+        dataset (:obj:`Dataset`): The dataset to be repeated.
+        times (int): Repeat times.
+    """
+
+    def __init__(self, dataset, times):
+        self.dataset = dataset
+        self.times = times
+        self.CLASSES = dataset.CLASSES
+
+        self._ori_len = len(self.dataset)
+
+    def __getitem__(self, idx):
+        return self.dataset[idx % self._ori_len]
+
+    def get_cat_ids(self, idx):
+        return self.dataset.get_cat_ids(idx % self._ori_len)
+
+    def __len__(self):
+        return self.times * self._ori_len
+
+
+# Modified from https://github.com/facebookresearch/detectron2/blob/41d475b75a230221e21d9cac5d69655e3415e3a4/detectron2/data/samplers/distributed_sampler.py#L57 # noqa
+@DATASETS.register_module()
+class ClassBalancedDataset(object):
+    """A wrapper of repeated dataset with repeat factor.
+
+    Suitable for training on class imbalanced datasets like LVIS. Following
+    the sampling strategy in [1], in each epoch, an image may appear multiple
+    times based on its "repeat factor".
+    The repeat factor for an image is a function of the frequency the rarest
+    category labeled in that image. The "frequency of category c" in [0, 1]
+    is defined by the fraction of images in the training set (without repeats)
+    in which category c appears.
+    The dataset needs to instantiate :func:`self.get_cat_ids(idx)` to support
+    ClassBalancedDataset.
+    The repeat factor is computed as followed.
+    1. For each category c, compute the fraction # of images
+        that contain it: f(c)
+    2. For each category c, compute the category-level repeat factor:
+        r(c) = max(1, sqrt(t/f(c)))
+    3. For each image I and its labels L(I), compute the image-level repeat
+    factor:
+        r(I) = max_{c in L(I)} r(c)
+
+    References:
+        .. [1]  https://arxiv.org/pdf/1908.03195.pdf
+
+    Args:
+        dataset (:obj:`CustomDataset`): The dataset to be repeated.
+        oversample_thr (float): frequency threshold below which data is
+            repeated. For categories with `f_c` >= `oversample_thr`, there is
+            no oversampling. For categories with `f_c` < `oversample_thr`, the
+            degree of oversampling following the square-root inverse frequency
+            heuristic above.
+    """
+
+    def __init__(self, dataset, oversample_thr, method='sqrt', **kwargs):
+        assert(method in ('sqrt', 'reciprocal'))    # reciprocal，倒数，绝对 balance， repeat_factor = f(c) / max( fc )
+        self.method = method
+        self.dataset = dataset
+        self.oversample_thr = oversample_thr
+        self.CLASSES = dataset.CLASSES
+
+        repeat_factors = self._get_repeat_factors(dataset, oversample_thr)
+        repeat_indices = []
+        for dataset_index, repeat_factor in enumerate(repeat_factors):
+            repeat_indices.extend([dataset_index] * math.ceil(repeat_factor))
+        self.repeat_indices = repeat_indices
+
+        flags = []
+        if hasattr(self.dataset, 'flag'):
+            for flag, repeat_factor in zip(self.dataset.flag, repeat_factors):
+                flags.extend([flag] * int(math.ceil(repeat_factor)))
+            assert len(flags) == len(repeat_indices)
+        self.flag = np.asarray(flags, dtype=np.uint8)
+
+    def _get_repeat_factors(self, dataset, repeat_thr):
+        # 1. For each category c, compute the fraction # of images
+        #   that contain it: f(c)
+        category_freq = defaultdict(int)
+        num_images = len(dataset)
+        for idx in range(num_images):
+            cat_ids = set(self.dataset.get_cat_ids(idx))
+            for cat_id in cat_ids:
+                category_freq[cat_id] += 1
+        for k, v in category_freq.items():
+            assert v > 0, f'caterogy {k} does not contain any images'
+            category_freq[k] = v / num_images
+
+        # 2. For each category c, compute the category-level repeat factor:
+        #    r(c) = max(1, sqrt(t/f(c)))
+        if self.method == 'sqrt':
+            category_repeat = {
+                cat_id: max(1.0, math.sqrt(repeat_thr / cat_freq))
+                for cat_id, cat_freq in category_freq.items()
+            }
+        elif self.method == 'reciprocal':
+            cat_freq_max = 0
+            for cat_id, cat_freq in category_freq.items():
+                cat_freq_max = max(cat_freq_max, cat_freq)
+            print('cat_freq_max: ', cat_freq_max)
+            category_repeat = {
+                cat_id: cat_freq_max / cat_freq
+                for cat_id, cat_freq in category_freq.items()
+            }
+        # 3. For each image I and its labels L(I), compute the image-level
+        # repeat factor:
+        #    r(I) = max_{c in L(I)} r(c)
+        repeat_factors = []
+        for idx in range(num_images):
+            cat_ids = set(self.dataset.get_cat_ids(idx))
+            repeat_factor = max(
+                {category_repeat[cat_id]
+                 for cat_id in cat_ids})
+            repeat_factors.append(repeat_factor)
+
+        return repeat_factors
+
+    def __getitem__(self, idx):
+        ori_index = self.repeat_indices[idx]
+        return self.dataset[ori_index]
+
+    def __len__(self):
+        return len(self.repeat_indices)

CAGE_expression_inference-apvit/apvit_mmcls/datasets/imagenet.py

@@ -0,0 +1,1105 @@
+import os
+
+import numpy as np
+
+from .base_dataset import BaseDataset
+from .builder import DATASETS
+
+
+def has_file_allowed_extension(filename, extensions):
+    """Checks if a file is an allowed extension.
+
+    Args:
+        filename (string): path to a file
+
+    Returns:
+        bool: True if the filename ends with a known image extension
+    """
+    filename_lower = filename.lower()
+    return any(filename_lower.endswith(ext) for ext in extensions)
+
+
+def find_folders(root):
+    """Find classes by folders under a root.
+
+    Args:
+        root (string): root directory of folders
+
+    Returns:
+        folder_to_idx (dict): the map from folder name to class idx
+    """
+    folders = [
+        d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))
+    ]
+    folders.sort()
+    folder_to_idx = {folders[i]: i for i in range(len(folders))}
+    return folder_to_idx
+
+
+def get_samples(root, folder_to_idx, extensions):
+    """Make dataset by walking all images under a root.
+
+    Args:
+        root (string): root directory of folders
+        folder_to_idx (dict): the map from class name to class idx
+        extensions (tuple): allowed extensions
+
+    Returns:
+        samples (list): a list of tuple where each element is (image, label)
+    """
+    samples = []
+    root = os.path.expanduser(root)
+    for folder_name in sorted(os.listdir(root)):
+        _dir = os.path.join(root, folder_name)
+        if not os.path.isdir(_dir):
+            continue
+
+        for _, _, fns in sorted(os.walk(_dir)):
+            for fn in sorted(fns):
+                if has_file_allowed_extension(fn, extensions):
+                    path = os.path.join(folder_name, fn)
+                    item = (path, folder_to_idx[folder_name])
+                    samples.append(item)
+    return samples
+
+
+@DATASETS.register_module()
+class ImageNet(BaseDataset):
+    """`ImageNet <http://www.image-net.org>`_ Dataset.
+
+    This implementation is modified from
+    https://github.com/pytorch/vision/blob/master/torchvision/datasets/imagenet.py  # noqa: E501
+    """
+
+    IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif')
+    CLASSES = [
+        'tench, Tinca tinca',
+        'goldfish, Carassius auratus',
+        'great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias',  # noqa: E501
+        'tiger shark, Galeocerdo cuvieri',
+        'hammerhead, hammerhead shark',
+        'electric ray, crampfish, numbfish, torpedo',
+        'stingray',
+        'cock',
+        'hen',
+        'ostrich, Struthio camelus',
+        'brambling, Fringilla montifringilla',
+        'goldfinch, Carduelis carduelis',
+        'house finch, linnet, Carpodacus mexicanus',
+        'junco, snowbird',
+        'indigo bunting, indigo finch, indigo bird, Passerina cyanea',
+        'robin, American robin, Turdus migratorius',
+        'bulbul',
+        'jay',
+        'magpie',
+        'chickadee',
+        'water ouzel, dipper',
+        'kite',
+        'bald eagle, American eagle, Haliaeetus leucocephalus',
+        'vulture',
+        'great grey owl, great gray owl, Strix nebulosa',
+        'European fire salamander, Salamandra salamandra',
+        'common newt, Triturus vulgaris',
+        'eft',
+        'spotted salamander, Ambystoma maculatum',
+        'axolotl, mud puppy, Ambystoma mexicanum',
+        'bullfrog, Rana catesbeiana',
+        'tree frog, tree-frog',
+        'tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui',
+        'loggerhead, loggerhead turtle, Caretta caretta',
+        'leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea',  # noqa: E501
+        'mud turtle',
+        'terrapin',
+        'box turtle, box tortoise',
+        'banded gecko',
+        'common iguana, iguana, Iguana iguana',
+        'American chameleon, anole, Anolis carolinensis',
+        'whiptail, whiptail lizard',
+        'agama',
+        'frilled lizard, Chlamydosaurus kingi',
+        'alligator lizard',
+        'Gila monster, Heloderma suspectum',
+        'green lizard, Lacerta viridis',
+        'African chameleon, Chamaeleo chamaeleon',
+        'Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis',  # noqa: E501
+        'African crocodile, Nile crocodile, Crocodylus niloticus',
+        'American alligator, Alligator mississipiensis',
+        'triceratops',
+        'thunder snake, worm snake, Carphophis amoenus',
+        'ringneck snake, ring-necked snake, ring snake',
+        'hognose snake, puff adder, sand viper',
+        'green snake, grass snake',
+        'king snake, kingsnake',
+        'garter snake, grass snake',
+        'water snake',
+        'vine snake',
+        'night snake, Hypsiglena torquata',
+        'boa constrictor, Constrictor constrictor',
+        'rock python, rock snake, Python sebae',
+        'Indian cobra, Naja naja',
+        'green mamba',
+        'sea snake',
+        'horned viper, cerastes, sand viper, horned asp, Cerastes cornutus',
+        'diamondback, diamondback rattlesnake, Crotalus adamanteus',
+        'sidewinder, horned rattlesnake, Crotalus cerastes',
+        'trilobite',
+        'harvestman, daddy longlegs, Phalangium opilio',
+        'scorpion',
+        'black and gold garden spider, Argiope aurantia',
+        'barn spider, Araneus cavaticus',
+        'garden spider, Aranea diademata',
+        'black widow, Latrodectus mactans',
+        'tarantula',
+        'wolf spider, hunting spider',
+        'tick',
+        'centipede',
+        'black grouse',
+        'ptarmigan',
+        'ruffed grouse, partridge, Bonasa umbellus',
+        'prairie chicken, prairie grouse, prairie fowl',
+        'peacock',
+        'quail',
+        'partridge',
+        'African grey, African gray, Psittacus erithacus',
+        'macaw',
+        'sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita',
+        'lorikeet',
+        'coucal',
+        'bee eater',
+        'hornbill',
+        'hummingbird',
+        'jacamar',
+        'toucan',
+        'drake',
+        'red-breasted merganser, Mergus serrator',
+        'goose',
+        'black swan, Cygnus atratus',
+        'tusker',
+        'echidna, spiny anteater, anteater',
+        'platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus',  # noqa: E501
+        'wallaby, brush kangaroo',
+        'koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus',  # noqa: E501
+        'wombat',
+        'jellyfish',
+        'sea anemone, anemone',
+        'brain coral',
+        'flatworm, platyhelminth',
+        'nematode, nematode worm, roundworm',
+        'conch',
+        'snail',
+        'slug',
+        'sea slug, nudibranch',
+        'chiton, coat-of-mail shell, sea cradle, polyplacophore',
+        'chambered nautilus, pearly nautilus, nautilus',
+        'Dungeness crab, Cancer magister',
+        'rock crab, Cancer irroratus',
+        'fiddler crab',
+        'king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica',  # noqa: E501
+        'American lobster, Northern lobster, Maine lobster, Homarus americanus',  # noqa: E501
+        'spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish',  # noqa: E501
+        'crayfish, crawfish, crawdad, crawdaddy',
+        'hermit crab',
+        'isopod',
+        'white stork, Ciconia ciconia',
+        'black stork, Ciconia nigra',
+        'spoonbill',
+        'flamingo',
+        'little blue heron, Egretta caerulea',
+        'American egret, great white heron, Egretta albus',
+        'bittern',
+        'crane',
+        'limpkin, Aramus pictus',
+        'European gallinule, Porphyrio porphyrio',
+        'American coot, marsh hen, mud hen, water hen, Fulica americana',
+        'bustard',
+        'ruddy turnstone, Arenaria interpres',
+        'red-backed sandpiper, dunlin, Erolia alpina',
+        'redshank, Tringa totanus',
+        'dowitcher',
+        'oystercatcher, oyster catcher',
+        'pelican',
+        'king penguin, Aptenodytes patagonica',
+        'albatross, mollymawk',
+        'grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus',  # noqa: E501
+        'killer whale, killer, orca, grampus, sea wolf, Orcinus orca',
+        'dugong, Dugong dugon',
+        'sea lion',
+        'Chihuahua',
+        'Japanese spaniel',
+        'Maltese dog, Maltese terrier, Maltese',
+        'Pekinese, Pekingese, Peke',
+        'Shih-Tzu',
+        'Blenheim spaniel',
+        'papillon',
+        'toy terrier',
+        'Rhodesian ridgeback',
+        'Afghan hound, Afghan',
+        'basset, basset hound',
+        'beagle',
+        'bloodhound, sleuthhound',
+        'bluetick',
+        'black-and-tan coonhound',
+        'Walker hound, Walker foxhound',
+        'English foxhound',
+        'redbone',
+        'borzoi, Russian wolfhound',
+        'Irish wolfhound',
+        'Italian greyhound',
+        'whippet',
+        'Ibizan hound, Ibizan Podenco',
+        'Norwegian elkhound, elkhound',
+        'otterhound, otter hound',
+        'Saluki, gazelle hound',
+        'Scottish deerhound, deerhound',
+        'Weimaraner',
+        'Staffordshire bullterrier, Staffordshire bull terrier',
+        'American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier',  # noqa: E501
+        'Bedlington terrier',
+        'Border terrier',
+        'Kerry blue terrier',
+        'Irish terrier',
+        'Norfolk terrier',
+        'Norwich terrier',
+        'Yorkshire terrier',
+        'wire-haired fox terrier',
+        'Lakeland terrier',
+        'Sealyham terrier, Sealyham',
+        'Airedale, Airedale terrier',
+        'cairn, cairn terrier',
+        'Australian terrier',
+        'Dandie Dinmont, Dandie Dinmont terrier',
+        'Boston bull, Boston terrier',
+        'miniature schnauzer',
+        'giant schnauzer',
+        'standard schnauzer',
+        'Scotch terrier, Scottish terrier, Scottie',
+        'Tibetan terrier, chrysanthemum dog',
+        'silky terrier, Sydney silky',
+        'soft-coated wheaten terrier',
+        'West Highland white terrier',
+        'Lhasa, Lhasa apso',
+        'flat-coated retriever',
+        'curly-coated retriever',
+        'golden retriever',
+        'Labrador retriever',
+        'Chesapeake Bay retriever',
+        'German short-haired pointer',
+        'vizsla, Hungarian pointer',
+        'English setter',
+        'Irish setter, red setter',
+        'Gordon setter',
+        'Brittany spaniel',
+        'clumber, clumber spaniel',
+        'English springer, English springer spaniel',
+        'Welsh springer spaniel',
+        'cocker spaniel, English cocker spaniel, cocker',
+        'Sussex spaniel',
+        'Irish water spaniel',
+        'kuvasz',
+        'schipperke',
+        'groenendael',
+        'malinois',
+        'briard',
+        'kelpie',
+        'komondor',
+        'Old English sheepdog, bobtail',
+        'Shetland sheepdog, Shetland sheep dog, Shetland',
+        'collie',
+        'Border collie',
+        'Bouvier des Flandres, Bouviers des Flandres',
+        'Rottweiler',
+        'German shepherd, German shepherd dog, German police dog, alsatian',
+        'Doberman, Doberman pinscher',
+        'miniature pinscher',
+        'Greater Swiss Mountain dog',
+        'Bernese mountain dog',
+        'Appenzeller',
+        'EntleBucher',
+        'boxer',
+        'bull mastiff',
+        'Tibetan mastiff',
+        'French bulldog',
+        'Great Dane',
+        'Saint Bernard, St Bernard',
+        'Eskimo dog, husky',
+        'malamute, malemute, Alaskan malamute',
+        'Siberian husky',
+        'dalmatian, coach dog, carriage dog',
+        'affenpinscher, monkey pinscher, monkey dog',
+        'basenji',
+        'pug, pug-dog',
+        'Leonberg',
+        'Newfoundland, Newfoundland dog',
+        'Great Pyrenees',
+        'Samoyed, Samoyede',
+        'Pomeranian',
+        'chow, chow chow',
+        'keeshond',
+        'Brabancon griffon',
+        'Pembroke, Pembroke Welsh corgi',
+        'Cardigan, Cardigan Welsh corgi',
+        'toy poodle',
+        'miniature poodle',
+        'standard poodle',
+        'Mexican hairless',
+        'timber wolf, grey wolf, gray wolf, Canis lupus',
+        'white wolf, Arctic wolf, Canis lupus tundrarum',
+        'red wolf, maned wolf, Canis rufus, Canis niger',
+        'coyote, prairie wolf, brush wolf, Canis latrans',
+        'dingo, warrigal, warragal, Canis dingo',
+        'dhole, Cuon alpinus',
+        'African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus',
+        'hyena, hyaena',
+        'red fox, Vulpes vulpes',
+        'kit fox, Vulpes macrotis',
+        'Arctic fox, white fox, Alopex lagopus',
+        'grey fox, gray fox, Urocyon cinereoargenteus',
+        'tabby, tabby cat',
+        'tiger cat',
+        'Persian cat',
+        'Siamese cat, Siamese',
+        'Egyptian cat',
+        'cougar, puma, catamount, mountain lion, painter, panther, Felis concolor',  # noqa: E501
+        'lynx, catamount',
+        'leopard, Panthera pardus',
+        'snow leopard, ounce, Panthera uncia',
+        'jaguar, panther, Panthera onca, Felis onca',
+        'lion, king of beasts, Panthera leo',
+        'tiger, Panthera tigris',
+        'cheetah, chetah, Acinonyx jubatus',
+        'brown bear, bruin, Ursus arctos',
+        'American black bear, black bear, Ursus americanus, Euarctos americanus',  # noqa: E501
+        'ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus',
+        'sloth bear, Melursus ursinus, Ursus ursinus',
+        'mongoose',
+        'meerkat, mierkat',
+        'tiger beetle',
+        'ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle',
+        'ground beetle, carabid beetle',
+        'long-horned beetle, longicorn, longicorn beetle',
+        'leaf beetle, chrysomelid',
+        'dung beetle',
+        'rhinoceros beetle',
+        'weevil',
+        'fly',
+        'bee',
+        'ant, emmet, pismire',
+        'grasshopper, hopper',
+        'cricket',
+        'walking stick, walkingstick, stick insect',
+        'cockroach, roach',
+        'mantis, mantid',
+        'cicada, cicala',
+        'leafhopper',
+        'lacewing, lacewing fly',
+        "dragonfly, darning needle, devil's darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk",  # noqa: E501
+        'damselfly',
+        'admiral',
+        'ringlet, ringlet butterfly',
+        'monarch, monarch butterfly, milkweed butterfly, Danaus plexippus',
+        'cabbage butterfly',
+        'sulphur butterfly, sulfur butterfly',
+        'lycaenid, lycaenid butterfly',
+        'starfish, sea star',
+        'sea urchin',
+        'sea cucumber, holothurian',
+        'wood rabbit, cottontail, cottontail rabbit',
+        'hare',
+        'Angora, Angora rabbit',
+        'hamster',
+        'porcupine, hedgehog',
+        'fox squirrel, eastern fox squirrel, Sciurus niger',
+        'marmot',
+        'beaver',
+        'guinea pig, Cavia cobaya',
+        'sorrel',
+        'zebra',
+        'hog, pig, grunter, squealer, Sus scrofa',
+        'wild boar, boar, Sus scrofa',
+        'warthog',
+        'hippopotamus, hippo, river horse, Hippopotamus amphibius',
+        'ox',
+        'water buffalo, water ox, Asiatic buffalo, Bubalus bubalis',
+        'bison',
+        'ram, tup',
+        'bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis',  # noqa: E501
+        'ibex, Capra ibex',
+        'hartebeest',
+        'impala, Aepyceros melampus',
+        'gazelle',
+        'Arabian camel, dromedary, Camelus dromedarius',
+        'llama',
+        'weasel',
+        'mink',
+        'polecat, fitch, foulmart, foumart, Mustela putorius',
+        'black-footed ferret, ferret, Mustela nigripes',
+        'otter',
+        'skunk, polecat, wood pussy',
+        'badger',
+        'armadillo',
+        'three-toed sloth, ai, Bradypus tridactylus',
+        'orangutan, orang, orangutang, Pongo pygmaeus',
+        'gorilla, Gorilla gorilla',
+        'chimpanzee, chimp, Pan troglodytes',
+        'gibbon, Hylobates lar',
+        'siamang, Hylobates syndactylus, Symphalangus syndactylus',
+        'guenon, guenon monkey',
+        'patas, hussar monkey, Erythrocebus patas',
+        'baboon',
+        'macaque',
+        'langur',
+        'colobus, colobus monkey',
+        'proboscis monkey, Nasalis larvatus',
+        'marmoset',
+        'capuchin, ringtail, Cebus capucinus',
+        'howler monkey, howler',
+        'titi, titi monkey',
+        'spider monkey, Ateles geoffroyi',
+        'squirrel monkey, Saimiri sciureus',
+        'Madagascar cat, ring-tailed lemur, Lemur catta',
+        'indri, indris, Indri indri, Indri brevicaudatus',
+        'Indian elephant, Elephas maximus',
+        'African elephant, Loxodonta africana',
+        'lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens',
+        'giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca',
+        'barracouta, snoek',
+        'eel',
+        'coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch',  # noqa: E501
+        'rock beauty, Holocanthus tricolor',
+        'anemone fish',
+        'sturgeon',
+        'gar, garfish, garpike, billfish, Lepisosteus osseus',
+        'lionfish',
+        'puffer, pufferfish, blowfish, globefish',
+        'abacus',
+        'abaya',
+        "academic gown, academic robe, judge's robe",
+        'accordion, piano accordion, squeeze box',
+        'acoustic guitar',
+        'aircraft carrier, carrier, flattop, attack aircraft carrier',
+        'airliner',
+        'airship, dirigible',
+        'altar',
+        'ambulance',
+        'amphibian, amphibious vehicle',
+        'analog clock',
+        'apiary, bee house',
+        'apron',
+        'ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin',  # noqa: E501
+        'assault rifle, assault gun',
+        'backpack, back pack, knapsack, packsack, rucksack, haversack',
+        'bakery, bakeshop, bakehouse',
+        'balance beam, beam',
+        'balloon',
+        'ballpoint, ballpoint pen, ballpen, Biro',
+        'Band Aid',
+        'banjo',
+        'bannister, banister, balustrade, balusters, handrail',
+        'barbell',
+        'barber chair',
+        'barbershop',
+        'barn',
+        'barometer',
+        'barrel, cask',
+        'barrow, garden cart, lawn cart, wheelbarrow',
+        'baseball',
+        'basketball',
+        'bassinet',
+        'bassoon',
+        'bathing cap, swimming cap',
+        'bath towel',
+        'bathtub, bathing tub, bath, tub',
+        'beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon',  # noqa: E501
+        'beacon, lighthouse, beacon light, pharos',
+        'beaker',
+        'bearskin, busby, shako',
+        'beer bottle',
+        'beer glass',
+        'bell cote, bell cot',
+        'bib',
+        'bicycle-built-for-two, tandem bicycle, tandem',
+        'bikini, two-piece',
+        'binder, ring-binder',
+        'binoculars, field glasses, opera glasses',
+        'birdhouse',
+        'boathouse',
+        'bobsled, bobsleigh, bob',
+        'bolo tie, bolo, bola tie, bola',
+        'bonnet, poke bonnet',
+        'bookcase',
+        'bookshop, bookstore, bookstall',
+        'bottlecap',
+        'bow',
+        'bow tie, bow-tie, bowtie',
+        'brass, memorial tablet, plaque',
+        'brassiere, bra, bandeau',
+        'breakwater, groin, groyne, mole, bulwark, seawall, jetty',
+        'breastplate, aegis, egis',
+        'broom',
+        'bucket, pail',
+        'buckle',
+        'bulletproof vest',
+        'bullet train, bullet',
+        'butcher shop, meat market',
+        'cab, hack, taxi, taxicab',
+        'caldron, cauldron',
+        'candle, taper, wax light',
+        'cannon',
+        'canoe',
+        'can opener, tin opener',
+        'cardigan',
+        'car mirror',
+        'carousel, carrousel, merry-go-round, roundabout, whirligig',
+        "carpenter's kit, tool kit",
+        'carton',
+        'car wheel',
+        'cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM',  # noqa: E501
+        'cassette',
+        'cassette player',
+        'castle',
+        'catamaran',
+        'CD player',
+        'cello, violoncello',
+        'cellular telephone, cellular phone, cellphone, cell, mobile phone',
+        'chain',
+        'chainlink fence',
+        'chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour',  # noqa: E501
+        'chain saw, chainsaw',
+        'chest',
+        'chiffonier, commode',
+        'chime, bell, gong',
+        'china cabinet, china closet',
+        'Christmas stocking',
+        'church, church building',
+        'cinema, movie theater, movie theatre, movie house, picture palace',
+        'cleaver, meat cleaver, chopper',
+        'cliff dwelling',
+        'cloak',
+        'clog, geta, patten, sabot',
+        'cocktail shaker',
+        'coffee mug',
+        'coffeepot',
+        'coil, spiral, volute, whorl, helix',
+        'combination lock',
+        'computer keyboard, keypad',
+        'confectionery, confectionary, candy store',
+        'container ship, containership, container vessel',
+        'convertible',
+        'corkscrew, bottle screw',
+        'cornet, horn, trumpet, trump',
+        'cowboy boot',
+        'cowboy hat, ten-gallon hat',
+        'cradle',
+        'crane',
+        'crash helmet',
+        'crate',
+        'crib, cot',
+        'Crock Pot',
+        'croquet ball',
+        'crutch',
+        'cuirass',
+        'dam, dike, dyke',
+        'desk',
+        'desktop computer',
+        'dial telephone, dial phone',
+        'diaper, nappy, napkin',
+        'digital clock',
+        'digital watch',
+        'dining table, board',
+        'dishrag, dishcloth',
+        'dishwasher, dish washer, dishwashing machine',
+        'disk brake, disc brake',
+        'dock, dockage, docking facility',
+        'dogsled, dog sled, dog sleigh',
+        'dome',
+        'doormat, welcome mat',
+        'drilling platform, offshore rig',
+        'drum, membranophone, tympan',
+        'drumstick',
+        'dumbbell',
+        'Dutch oven',
+        'electric fan, blower',
+        'electric guitar',
+        'electric locomotive',
+        'entertainment center',
+        'envelope',
+        'espresso maker',
+        'face powder',
+        'feather boa, boa',
+        'file, file cabinet, filing cabinet',
+        'fireboat',
+        'fire engine, fire truck',
+        'fire screen, fireguard',
+        'flagpole, flagstaff',
+        'flute, transverse flute',
+        'folding chair',
+        'football helmet',
+        'forklift',
+        'fountain',
+        'fountain pen',
+        'four-poster',
+        'freight car',
+        'French horn, horn',
+        'frying pan, frypan, skillet',
+        'fur coat',
+        'garbage truck, dustcart',
+        'gasmask, respirator, gas helmet',
+        'gas pump, gasoline pump, petrol pump, island dispenser',
+        'goblet',
+        'go-kart',
+        'golf ball',
+        'golfcart, golf cart',
+        'gondola',
+        'gong, tam-tam',
+        'gown',
+        'grand piano, grand',
+        'greenhouse, nursery, glasshouse',
+        'grille, radiator grille',
+        'grocery store, grocery, food market, market',
+        'guillotine',
+        'hair slide',
+        'hair spray',
+        'half track',
+        'hammer',
+        'hamper',
+        'hand blower, blow dryer, blow drier, hair dryer, hair drier',
+        'hand-held computer, hand-held microcomputer',
+        'handkerchief, hankie, hanky, hankey',
+        'hard disc, hard disk, fixed disk',
+        'harmonica, mouth organ, harp, mouth harp',
+        'harp',
+        'harvester, reaper',
+        'hatchet',
+        'holster',
+        'home theater, home theatre',
+        'honeycomb',
+        'hook, claw',
+        'hoopskirt, crinoline',
+        'horizontal bar, high bar',
+        'horse cart, horse-cart',
+        'hourglass',
+        'iPod',
+        'iron, smoothing iron',
+        "jack-o'-lantern",
+        'jean, blue jean, denim',
+        'jeep, landrover',
+        'jersey, T-shirt, tee shirt',
+        'jigsaw puzzle',
+        'jinrikisha, ricksha, rickshaw',
+        'joystick',
+        'kimono',
+        'knee pad',
+        'knot',
+        'lab coat, laboratory coat',
+        'ladle',
+        'lampshade, lamp shade',
+        'laptop, laptop computer',
+        'lawn mower, mower',
+        'lens cap, lens cover',
+        'letter opener, paper knife, paperknife',
+        'library',
+        'lifeboat',
+        'lighter, light, igniter, ignitor',
+        'limousine, limo',
+        'liner, ocean liner',
+        'lipstick, lip rouge',
+        'Loafer',
+        'lotion',
+        'loudspeaker, speaker, speaker unit, loudspeaker system, speaker system',  # noqa: E501
+        "loupe, jeweler's loupe",
+        'lumbermill, sawmill',
+        'magnetic compass',
+        'mailbag, postbag',
+        'mailbox, letter box',
+        'maillot',
+        'maillot, tank suit',
+        'manhole cover',
+        'maraca',
+        'marimba, xylophone',
+        'mask',
+        'matchstick',
+        'maypole',
+        'maze, labyrinth',
+        'measuring cup',
+        'medicine chest, medicine cabinet',
+        'megalith, megalithic structure',
+        'microphone, mike',
+        'microwave, microwave oven',
+        'military uniform',
+        'milk can',
+        'minibus',
+        'miniskirt, mini',
+        'minivan',
+        'missile',
+        'mitten',
+        'mixing bowl',
+        'mobile home, manufactured home',
+        'Model T',
+        'modem',
+        'monastery',
+        'monitor',
+        'moped',
+        'mortar',
+        'mortarboard',
+        'mosque',
+        'mosquito net',
+        'motor scooter, scooter',
+        'mountain bike, all-terrain bike, off-roader',
+        'mountain tent',
+        'mouse, computer mouse',
+        'mousetrap',
+        'moving van',
+        'muzzle',
+        'nail',
+        'neck brace',
+        'necklace',
+        'nipple',
+        'notebook, notebook computer',
+        'obelisk',
+        'oboe, hautboy, hautbois',
+        'ocarina, sweet potato',
+        'odometer, hodometer, mileometer, milometer',
+        'oil filter',
+        'organ, pipe organ',
+        'oscilloscope, scope, cathode-ray oscilloscope, CRO',
+        'overskirt',
+        'oxcart',
+        'oxygen mask',
+        'packet',
+        'paddle, boat paddle',
+        'paddlewheel, paddle wheel',
+        'padlock',
+        'paintbrush',
+        "pajama, pyjama, pj's, jammies",
+        'palace',
+        'panpipe, pandean pipe, syrinx',
+        'paper towel',
+        'parachute, chute',
+        'parallel bars, bars',
+        'park bench',
+        'parking meter',
+        'passenger car, coach, carriage',
+        'patio, terrace',
+        'pay-phone, pay-station',
+        'pedestal, plinth, footstall',
+        'pencil box, pencil case',
+        'pencil sharpener',
+        'perfume, essence',
+        'Petri dish',
+        'photocopier',
+        'pick, plectrum, plectron',
+        'pickelhaube',
+        'picket fence, paling',
+        'pickup, pickup truck',
+        'pier',
+        'piggy bank, penny bank',
+        'pill bottle',
+        'pillow',
+        'ping-pong ball',
+        'pinwheel',
+        'pirate, pirate ship',
+        'pitcher, ewer',
+        "plane, carpenter's plane, woodworking plane",
+        'planetarium',
+        'plastic bag',
+        'plate rack',
+        'plow, plough',
+        "plunger, plumber's helper",
+        'Polaroid camera, Polaroid Land camera',
+        'pole',
+        'police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria',  # noqa: E501
+        'poncho',
+        'pool table, billiard table, snooker table',
+        'pop bottle, soda bottle',
+        'pot, flowerpot',
+        "potter's wheel",
+        'power drill',
+        'prayer rug, prayer mat',
+        'printer',
+        'prison, prison house',
+        'projectile, missile',
+        'projector',
+        'puck, hockey puck',
+        'punching bag, punch bag, punching ball, punchball',
+        'purse',
+        'quill, quill pen',
+        'quilt, comforter, comfort, puff',
+        'racer, race car, racing car',
+        'racket, racquet',
+        'radiator',
+        'radio, wireless',
+        'radio telescope, radio reflector',
+        'rain barrel',
+        'recreational vehicle, RV, R.V.',
+        'reel',
+        'reflex camera',
+        'refrigerator, icebox',
+        'remote control, remote',
+        'restaurant, eating house, eating place, eatery',
+        'revolver, six-gun, six-shooter',
+        'rifle',
+        'rocking chair, rocker',
+        'rotisserie',
+        'rubber eraser, rubber, pencil eraser',
+        'rugby ball',
+        'rule, ruler',
+        'running shoe',
+        'safe',
+        'safety pin',
+        'saltshaker, salt shaker',
+        'sandal',
+        'sarong',
+        'sax, saxophone',
+        'scabbard',
+        'scale, weighing machine',
+        'school bus',
+        'schooner',
+        'scoreboard',
+        'screen, CRT screen',
+        'screw',
+        'screwdriver',
+        'seat belt, seatbelt',
+        'sewing machine',
+        'shield, buckler',
+        'shoe shop, shoe-shop, shoe store',
+        'shoji',
+        'shopping basket',
+        'shopping cart',
+        'shovel',
+        'shower cap',
+        'shower curtain',
+        'ski',
+        'ski mask',
+        'sleeping bag',
+        'slide rule, slipstick',
+        'sliding door',
+        'slot, one-armed bandit',
+        'snorkel',
+        'snowmobile',
+        'snowplow, snowplough',
+        'soap dispenser',
+        'soccer ball',
+        'sock',
+        'solar dish, solar collector, solar furnace',
+        'sombrero',
+        'soup bowl',
+        'space bar',
+        'space heater',
+        'space shuttle',
+        'spatula',
+        'speedboat',
+        "spider web, spider's web",
+        'spindle',
+        'sports car, sport car',
+        'spotlight, spot',
+        'stage',
+        'steam locomotive',
+        'steel arch bridge',
+        'steel drum',
+        'stethoscope',
+        'stole',
+        'stone wall',
+        'stopwatch, stop watch',
+        'stove',
+        'strainer',
+        'streetcar, tram, tramcar, trolley, trolley car',
+        'stretcher',
+        'studio couch, day bed',
+        'stupa, tope',
+        'submarine, pigboat, sub, U-boat',
+        'suit, suit of clothes',
+        'sundial',
+        'sunglass',
+        'sunglasses, dark glasses, shades',
+        'sunscreen, sunblock, sun blocker',
+        'suspension bridge',
+        'swab, swob, mop',
+        'sweatshirt',
+        'swimming trunks, bathing trunks',
+        'swing',
+        'switch, electric switch, electrical switch',
+        'syringe',
+        'table lamp',
+        'tank, army tank, armored combat vehicle, armoured combat vehicle',
+        'tape player',
+        'teapot',
+        'teddy, teddy bear',
+        'television, television system',
+        'tennis ball',
+        'thatch, thatched roof',
+        'theater curtain, theatre curtain',
+        'thimble',
+        'thresher, thrasher, threshing machine',
+        'throne',
+        'tile roof',
+        'toaster',
+        'tobacco shop, tobacconist shop, tobacconist',
+        'toilet seat',
+        'torch',
+        'totem pole',
+        'tow truck, tow car, wrecker',
+        'toyshop',
+        'tractor',
+        'trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi',  # noqa: E501
+        'tray',
+        'trench coat',
+        'tricycle, trike, velocipede',
+        'trimaran',
+        'tripod',
+        'triumphal arch',
+        'trolleybus, trolley coach, trackless trolley',
+        'trombone',
+        'tub, vat',
+        'turnstile',
+        'typewriter keyboard',
+        'umbrella',
+        'unicycle, monocycle',
+        'upright, upright piano',
+        'vacuum, vacuum cleaner',
+        'vase',
+        'vault',
+        'velvet',
+        'vending machine',
+        'vestment',
+        'viaduct',
+        'violin, fiddle',
+        'volleyball',
+        'waffle iron',
+        'wall clock',
+        'wallet, billfold, notecase, pocketbook',
+        'wardrobe, closet, press',
+        'warplane, military plane',
+        'washbasin, handbasin, washbowl, lavabo, wash-hand basin',
+        'washer, automatic washer, washing machine',
+        'water bottle',
+        'water jug',
+        'water tower',
+        'whiskey jug',
+        'whistle',
+        'wig',
+        'window screen',
+        'window shade',
+        'Windsor tie',
+        'wine bottle',
+        'wing',
+        'wok',
+        'wooden spoon',
+        'wool, woolen, woollen',
+        'worm fence, snake fence, snake-rail fence, Virginia fence',
+        'wreck',
+        'yawl',
+        'yurt',
+        'web site, website, internet site, site',
+        'comic book',
+        'crossword puzzle, crossword',
+        'street sign',
+        'traffic light, traffic signal, stoplight',
+        'book jacket, dust cover, dust jacket, dust wrapper',
+        'menu',
+        'plate',
+        'guacamole',
+        'consomme',
+        'hot pot, hotpot',
+        'trifle',
+        'ice cream, icecream',
+        'ice lolly, lolly, lollipop, popsicle',
+        'French loaf',
+        'bagel, beigel',
+        'pretzel',
+        'cheeseburger',
+        'hotdog, hot dog, red hot',
+        'mashed potato',
+        'head cabbage',
+        'broccoli',
+        'cauliflower',
+        'zucchini, courgette',
+        'spaghetti squash',
+        'acorn squash',
+        'butternut squash',
+        'cucumber, cuke',
+        'artichoke, globe artichoke',
+        'bell pepper',
+        'cardoon',
+        'mushroom',
+        'Granny Smith',
+        'strawberry',
+        'orange',
+        'lemon',
+        'fig',
+        'pineapple, ananas',
+        'banana',
+        'jackfruit, jak, jack',
+        'custard apple',
+        'pomegranate',
+        'hay',
+        'carbonara',
+        'chocolate sauce, chocolate syrup',
+        'dough',
+        'meat loaf, meatloaf',
+        'pizza, pizza pie',
+        'potpie',
+        'burrito',
+        'red wine',
+        'espresso',
+        'cup',
+        'eggnog',
+        'alp',
+        'bubble',
+        'cliff, drop, drop-off',
+        'coral reef',
+        'geyser',
+        'lakeside, lakeshore',
+        'promontory, headland, head, foreland',
+        'sandbar, sand bar',
+        'seashore, coast, seacoast, sea-coast',
+        'valley, vale',
+        'volcano',
+        'ballplayer, baseball player',
+        'groom, bridegroom',
+        'scuba diver',
+        'rapeseed',
+        'daisy',
+        "yellow lady's slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum",  # noqa: E501
+        'corn',
+        'acorn',
+        'hip, rose hip, rosehip',
+        'buckeye, horse chestnut, conker',
+        'coral fungus',
+        'agaric',
+        'gyromitra',
+        'stinkhorn, carrion fungus',
+        'earthstar',
+        'hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa',  # noqa: E501
+        'bolete',
+        'ear, spike, capitulum',
+        'toilet tissue, toilet paper, bathroom tissue'
+    ]
+
+    def load_annotations(self):
+        if self.ann_file is None:
+            folder_to_idx = find_folders(self.data_prefix)
+            samples = get_samples(
+                self.data_prefix,
+                folder_to_idx,
+                extensions=self.IMG_EXTENSIONS)
+            if len(samples) == 0:
+                raise (RuntimeError('Found 0 files in subfolders of: '
+                                    f'{self.data_prefix}. '
+                                    'Supported extensions are: '
+                                    f'{",".join(self.IMG_EXTENSIONS)}'))
+
+            self.folder_to_idx = folder_to_idx
+        elif isinstance(self.ann_file, str):
+            with open(self.ann_file) as f:
+                samples = [x.strip().split(' ') for x in f.readlines()]
+        else:
+            raise TypeError('ann_file must be a str or None')
+        self.samples = samples
+
+        data_infos = []
+        for filename, gt_label in self.samples:
+            info = {'img_prefix': self.data_prefix}
+            info['img_info'] = {'filename': filename}
+            info['gt_label'] = np.array(gt_label, dtype=np.int64)
+            data_infos.append(info)
+        return data_infos

CAGE_expression_inference-apvit/apvit_mmcls/datasets/mnist.py

@@ -0,0 +1,171 @@
+import codecs
+import os
+import os.path as osp
+
+import numpy as np
+import torch
+
+from .base_dataset import BaseDataset
+from .builder import DATASETS
+from .utils import download_and_extract_archive, rm_suffix
+
+
+@DATASETS.register_module()
+class MNIST(BaseDataset):
+    """`MNIST <http://yann.lecun.com/exdb/mnist/>`_ Dataset.
+
+    This implementation is modified from
+    https://github.com/pytorch/vision/blob/master/torchvision/datasets/mnist.py  # noqa: E501
+    """
+
+    resource_prefix = 'http://yann.lecun.com/exdb/mnist/'
+    resources = {
+        'train_image_file':
+        ('train-images-idx3-ubyte.gz', 'f68b3c2dcbeaaa9fbdd348bbdeb94873'),
+        'train_label_file':
+        ('train-labels-idx1-ubyte.gz', 'd53e105ee54ea40749a09fcbcd1e9432'),
+        'test_image_file':
+        ('t10k-images-idx3-ubyte.gz', '9fb629c4189551a2d022fa330f9573f3'),
+        'test_label_file':
+        ('t10k-labels-idx1-ubyte.gz', 'ec29112dd5afa0611ce80d1b7f02629c')
+    }
+
+    CLASSES = [
+        '0 - zero', '1 - one', '2 - two', '3 - three', '4 - four', '5 - five',
+        '6 - six', '7 - seven', '8 - eight', '9 - nine'
+    ]
+
+    def load_annotations(self):
+        train_image_file = osp.join(
+            self.data_prefix, rm_suffix(self.resources['train_image_file'][0]))
+        train_label_file = osp.join(
+            self.data_prefix, rm_suffix(self.resources['train_label_file'][0]))
+        test_image_file = osp.join(
+            self.data_prefix, rm_suffix(self.resources['test_image_file'][0]))
+        test_label_file = osp.join(
+            self.data_prefix, rm_suffix(self.resources['test_label_file'][0]))
+
+        if not osp.exists(train_image_file) or not osp.exists(
+                train_label_file) or not osp.exists(
+                    test_image_file) or not osp.exists(test_label_file):
+            self.download()
+
+        train_set = (read_image_file(train_image_file),
+                     read_label_file(train_label_file))
+        test_set = (read_image_file(test_image_file),
+                    read_label_file(test_label_file))
+
+        if not self.test_mode:
+            imgs, gt_labels = train_set
+        else:
+            imgs, gt_labels = test_set
+
+        data_infos = []
+        for img, gt_label in zip(imgs, gt_labels):
+            gt_label = np.array(gt_label, dtype=np.int64)
+            info = {'img': img.numpy(), 'gt_label': gt_label}
+            data_infos.append(info)
+        return data_infos
+
+    def download(self):
+        os.makedirs(self.data_prefix, exist_ok=True)
+
+        # download files
+        for url, md5 in self.resources.values():
+            url = osp.join(self.resource_prefix, url)
+            filename = url.rpartition('/')[2]
+            download_and_extract_archive(
+                url,
+                download_root=self.data_prefix,
+                filename=filename,
+                md5=md5)
+
+
+@DATASETS.register_module()
+class FashionMNIST(MNIST):
+    """`Fashion-MNIST <https://github.com/zalandoresearch/fashion-mnist>`_
+        Dataset.
+    """
+
+    resource_prefix = 'http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/'  # noqa: E501
+    resources = {
+        'train_image_file':
+        ('train-images-idx3-ubyte.gz', '8d4fb7e6c68d591d4c3dfef9ec88bf0d'),
+        'train_label_file':
+        ('train-labels-idx1-ubyte.gz', '25c81989df183df01b3e8a0aad5dffbe'),
+        'test_image_file':
+        ('t10k-images-idx3-ubyte.gz', 'bef4ecab320f06d8554ea6380940ec79'),
+        'test_label_file':
+        ('t10k-labels-idx1-ubyte.gz', 'bb300cfdad3c16e7a12a480ee83cd310')
+    }
+    CLASSES = [
+        'T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal',
+        'Shirt', 'Sneaker', 'Bag', 'Ankle boot'
+    ]
+
+
+def get_int(b):
+    return int(codecs.encode(b, 'hex'), 16)
+
+
+def open_maybe_compressed_file(path):
+    """Return a file object that possibly decompresses 'path' on the fly.
+       Decompression occurs when argument `path` is a string
+       and ends with '.gz' or '.xz'.
+    """
+    if not isinstance(path, str):
+        return path
+    if path.endswith('.gz'):
+        import gzip
+        return gzip.open(path, 'rb')
+    if path.endswith('.xz'):
+        import lzma
+        return lzma.open(path, 'rb')
+    return open(path, 'rb')
+
+
+def read_sn3_pascalvincent_tensor(path, strict=True):
+    """Read a SN3 file in "Pascal Vincent" format
+       (Lush file 'libidx/idx-io.lsh').
+       Argument may be a filename, compressed filename, or file object.
+    """
+    # typemap
+    if not hasattr(read_sn3_pascalvincent_tensor, 'typemap'):
+        read_sn3_pascalvincent_tensor.typemap = {
+            8: (torch.uint8, np.uint8, np.uint8),
+            9: (torch.int8, np.int8, np.int8),
+            11: (torch.int16, np.dtype('>i2'), 'i2'),
+            12: (torch.int32, np.dtype('>i4'), 'i4'),
+            13: (torch.float32, np.dtype('>f4'), 'f4'),
+            14: (torch.float64, np.dtype('>f8'), 'f8')
+        }
+    # read
+    with open_maybe_compressed_file(path) as f:
+        data = f.read()
+    # parse
+    magic = get_int(data[0:4])
+    nd = magic % 256
+    ty = magic // 256
+    assert nd >= 1 and nd <= 3
+    assert ty >= 8 and ty <= 14
+    m = read_sn3_pascalvincent_tensor.typemap[ty]
+    s = [get_int(data[4 * (i + 1):4 * (i + 2)]) for i in range(nd)]
+    parsed = np.frombuffer(data, dtype=m[1], offset=(4 * (nd + 1)))
+    assert parsed.shape[0] == np.prod(s) or not strict
+    return torch.from_numpy(parsed.astype(m[2], copy=False)).view(*s)
+
+
+def read_label_file(path):
+    with open(path, 'rb') as f:
+        x = read_sn3_pascalvincent_tensor(f, strict=False)
+    assert (x.dtype == torch.uint8)
+    assert (x.ndimension() == 1)
+    return x.long()
+
+
+def read_image_file(path):
+    with open(path, 'rb') as f:
+        x = read_sn3_pascalvincent_tensor(f, strict=False)
+    assert (x.dtype == torch.uint8)
+    assert (x.ndimension() == 3)
+    return x

CAGE_expression_inference-apvit/apvit_mmcls/datasets/pipelines/__init__.py

@@ -0,0 +1,15 @@
+from .compose import Compose
+from .formating import (Collect, ImageToTensor, ToNumpy, ToPIL, ToTensor,
+                        Transpose, to_tensor)
+from .loading import LoadImageFromFile
+from .transforms import (RandomAppliedTrans, CenterCrop, RandomCrop, RandomFlip, RandomGrayscale,
+                         RandomResizedCrop, Resize, RandomRotate, ColorJitter)
+from .test_time_aug import MultiScaleFlipAug
+
+__all__ = [
+    'Compose', 'to_tensor', 'ToTensor', 'ImageToTensor', 'ToPIL', 'ToNumpy',
+    'Transpose', 'Collect', 'LoadImageFromFile', 'Resize', 'CenterCrop',
+    'RandomFlip', 'Normalize', 'RandomCrop', 'RandomResizedCrop',
+    'RandomGrayscale', 'RandomAppliedTrans', 'RandomRotate', 'ColorJitter',
+    'MultiScaleFlipAug'
+]

CAGE_expression_inference-apvit/apvit_mmcls/datasets/pipelines/compose.py

@@ -0,0 +1,42 @@
+from collections.abc import Sequence
+
+from mmcv.utils import build_from_cfg
+
+from ..builder import PIPELINES
+
+
+@PIPELINES.register_module()
+class Compose(object):
+    """Compose a data pipeline with a sequence of transforms.
+
+    Args:
+        transforms (list[dict | callable]):
+            Either config dicts of transforms or transform objects.
+    """
+
+    def __init__(self, transforms):
+        assert isinstance(transforms, Sequence)
+        self.transforms = []
+        for transform in transforms:
+            if isinstance(transform, dict):
+                transform = build_from_cfg(transform, PIPELINES)
+                self.transforms.append(transform)
+            elif callable(transform):
+                self.transforms.append(transform)
+            else:
+                raise TypeError('transform must be callable or a dict, but got'
+                                f' {type(transform)}')
+
+    def __call__(self, data):
+        for t in self.transforms:
+            data = t(data)
+            if data is None:
+                return None
+        return data
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '('
+        for t in self.transforms:
+            format_string += f'\n    {t}'
+        format_string += '\n)'
+        return format_string

CAGE_expression_inference-apvit/apvit_mmcls/datasets/pipelines/formating.py

@@ -0,0 +1,156 @@
+from collections.abc import Sequence
+
+import mmcv
+import numpy as np
+import torch
+from PIL import Image
+
+from ..builder import PIPELINES
+
+
+def to_tensor(data):
+    """Convert objects of various python types to :obj:`torch.Tensor`.
+
+    Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`,
+    :class:`Sequence`, :class:`int` and :class:`float`.
+    """
+    if isinstance(data, torch.Tensor):
+        return data
+    elif isinstance(data, np.ndarray):
+        return torch.from_numpy(data)
+    elif isinstance(data, Sequence) and not mmcv.is_str(data):
+        return torch.tensor(data)
+    elif isinstance(data, int):
+        return torch.LongTensor([data])
+    elif isinstance(data, float):
+        return torch.FloatTensor([data])
+    else:
+        raise TypeError(
+            f'Type {type(data)} cannot be converted to tensor.'
+            'Supported types are: `numpy.ndarray`, `torch.Tensor`, '
+            '`Sequence`, `int` and `float`')
+
+
+@PIPELINES.register_module()
+class ToTensor(object):
+
+    def __init__(self, keys):
+        self.keys = keys
+
+    def __call__(self, results):
+        for key in self.keys:
+            results[key] = to_tensor(results[key])
+        return results
+
+    def __repr__(self):
+        return self.__class__.__name__ + f'(keys={self.keys})'
+
+
+@PIPELINES.register_module()
+class ImageToTensor(object):
+
+    def __init__(self, keys):
+        self.keys = keys
+
+    def __call__(self, results):
+        for key in self.keys:
+            img = results[key]
+            if len(img.shape) < 3:
+                img = np.expand_dims(img, -1)
+            results[key] = to_tensor(img.transpose(2, 0, 1))
+        return results
+
+    def __repr__(self):
+        return self.__class__.__name__ + f'(keys={self.keys})'
+
+
+@PIPELINES.register_module()
+class Transpose(object):
+
+    def __init__(self, keys, order):
+        self.keys = keys
+        self.order = order
+
+    def __call__(self, results):
+        for key in self.keys:
+            results[key] = results[key].transpose(self.order)
+        return results
+
+    def __repr__(self):
+        return self.__class__.__name__ + \
+            f'(keys={self.keys}, order={self.order})'
+
+
+@PIPELINES.register_module()
+class ToPIL(object):
+
+    def __init__(self):
+        pass
+
+    def __call__(self, results):
+        results['img'] = Image.fromarray(results['img'])
+        return results
+
+
+@PIPELINES.register_module()
+class ToNumpy(object):
+
+    def __init__(self):
+        pass
+
+    def __call__(self, results):
+        results['img'] = np.array(results['img'], dtype=np.float32)
+        return results
+
+
+@PIPELINES.register_module()
+class Collect(object):
+    """
+    Collect data from the loader relevant to the specific task.
+
+    This is usually the last stage of the data loader pipeline. Typically keys
+    is set to some subset of "img" and "gt_label".
+    """
+
+    def __init__(self, keys):
+        self.keys = keys
+
+    def __call__(self, results):
+        data = {}
+        for key in self.keys:
+            data[key] = results[key]
+        return data
+
+    def __repr__(self):
+        return self.__class__.__name__ + \
+            f'(keys={self.keys}, meta_keys={self.meta_keys})'
+
+
+@PIPELINES.register_module()
+class WrapFieldsToLists(object):
+    """Wrap fields of the data dictionary into lists for evaluation.
+
+    This class can be used as a last step of a test or validation
+    pipeline for single image evaluation or inference.
+
+    Example:
+        >>> test_pipeline = [
+        >>>    dict(type='LoadImageFromFile'),
+        >>>    dict(type='Normalize',
+                    mean=[123.675, 116.28, 103.53],
+                    std=[58.395, 57.12, 57.375],
+                    to_rgb=True),
+        >>>    dict(type='ImageToTensor', keys=['img']),
+        >>>    dict(type='Collect', keys=['img']),
+        >>>    dict(type='WrapIntoLists')
+        >>> ]
+    """
+
+    def __call__(self, results):
+        # Wrap dict fields into lists
+        for key, val in results.items():
+            results[key] = [val]
+        return results
+
+    def __repr__(self):
+        return f'{self.__class__.__name__}()'

CAGE_expression_inference-apvit/apvit_mmcls/datasets/pipelines/loading.py

@@ -0,0 +1,67 @@
+import os.path as osp
+
+import mmcv
+import numpy as np
+
+from ..builder import PIPELINES
+
+
+@PIPELINES.register_module()
+class LoadImageFromFile(object):
+    """Load an image from file.
+
+    Required keys are "img_prefix" and "img_info" (a dict that must contain the
+    key "filename"). Added or updated keys are "filename", "img", "img_shape",
+    "ori_shape" (same as `img_shape`) and "img_norm_cfg" (means=0 and stds=1).
+
+    Args:
+        to_float32 (bool): Whether to convert the loaded image to a float32
+            numpy array. If set to False, the loaded image is an uint8 array.
+            Defaults to False.
+        color_type (str): The flag argument for :func:`mmcv.imfrombytes()`.
+            Defaults to 'color'.
+        file_client_args (dict): Arguments to instantiate a FileClient.
+            See :class:`mmcv.fileio.FileClient` for details.
+            Defaults to ``dict(backend='disk')``.
+    """
+
+    def __init__(self,
+                 to_float32=False,
+                 color_type='color',
+                 file_client_args=dict(backend='disk')):
+        self.to_float32 = to_float32
+        self.color_type = color_type
+        self.file_client_args = file_client_args.copy()
+        self.file_client = None
+
+    def __call__(self, results):
+        if self.file_client is None:
+            self.file_client = mmcv.FileClient(**self.file_client_args)
+
+        if results['img_prefix'] is not None:
+            filename = osp.join(results['img_prefix'],
+                                results['img_info']['filename'])
+        else:
+            filename = results['img_info']['filename']
+
+        img_bytes = self.file_client.get(filename)
+        img = mmcv.imfrombytes(img_bytes, flag=self.color_type)
+        if self.to_float32:
+            img = img.astype(np.float32)
+        results['filename'] = filename
+        results['img'] = img
+        results['img_shape'] = img.shape
+        results['ori_shape'] = img.shape
+        num_channels = 1 if len(img.shape) < 3 else img.shape[2]
+        results['img_norm_cfg'] = dict(
+            mean=np.zeros(num_channels, dtype=np.float32),
+            std=np.ones(num_channels, dtype=np.float32),
+            to_rgb=False)
+        return results
+
+    def __repr__(self):
+        repr_str = (f'{self.__class__.__name__}('
+                    f'to_float32={self.to_float32}, '
+                    f"color_type='{self.color_type}', "
+                    f'file_client_args={self.file_client_args})')
+        return repr_str

CAGE_expression_inference-apvit/apvit_mmcls/datasets/pipelines/test_time_aug.py

@@ -0,0 +1,126 @@
+import warnings
+
+import mmcv
+
+from ..builder import PIPELINES
+from .compose import Compose
+
+
+@PIPELINES.register_module()
+class MultiScaleFlipAug(object):
+    """Test-time augmentation with multiple scales and flipping.
+
+    An example configuration is as followed:
+
+    .. code-block::
+
+        img_scale=(2048, 1024),
+        img_ratios=[0.5, 1.0],
+        flip=True,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ]
+
+    After MultiScaleFLipAug with above configuration, the results are wrapped
+    into lists of the same length as followed:
+
+    .. code-block::
+
+        dict(
+            img=[...],
+            img_shape=[...],
+            scale=[(1024, 512), (1024, 512), (2048, 1024), (2048, 1024)]
+            flip=[False, True, False, True]
+            ...
+        )
+
+    Args:
+        transforms (list[dict]): Transforms to apply in each augmentation.
+        img_scale (tuple | list[tuple]): Images scales for resizing.
+        img_ratios (float | list[float]): Image ratios for resizing
+        flip (bool): Whether apply flip augmentation. Default: False.
+        flip_direction (str | list[str]): Flip augmentation directions,
+            options are "horizontal" and "vertical". If flip_direction is list,
+            multiple flip augmentations will be applied.
+            It has no effect when flip == False. Default: "horizontal".
+    """
+
+    def __init__(self,
+                 transforms,
+                 img_scale,
+                 num=8,
+                 img_ratios=None,
+                 flip=False,
+                 flip_direction='horizontal'):
+        self.transforms = Compose(transforms)
+        self.num = num
+        if img_ratios is not None:
+            # mode 1: given a scale and a range of image ratio
+            img_ratios = img_ratios if isinstance(img_ratios,
+                                                  list) else [img_ratios]
+            assert mmcv.is_list_of(img_ratios, float)
+            assert isinstance(img_scale, tuple) and len(img_scale) == 2
+            self.img_scale = [(int(img_scale[0] * ratio),
+                               int(img_scale[1] * ratio))
+                              for ratio in img_ratios]
+        else:
+            # mode 2: given multiple scales
+            self.img_scale = img_scale if isinstance(img_scale,
+                                                     list) else [img_scale]
+        assert mmcv.is_list_of(self.img_scale, tuple)
+        self.flip = flip
+        self.flip_direction = flip_direction if isinstance(
+            flip_direction, list) else [flip_direction]
+        assert mmcv.is_list_of(self.flip_direction, str)
+        if not self.flip and self.flip_direction != ['horizontal']:
+            warnings.warn(
+                'flip_direction has no effect when flip is set to False')
+        if (self.flip
+                and not any([t['type'] == 'RandomFlip' for t in transforms])):
+            warnings.warn(
+                'flip has no effect when RandomFlip is not in transforms')
+
+    def __call__(self, results):
+        """Call function to apply test time augment transforms on results.
+
+        Args:
+            results (dict): Result dict contains the data to transform.
+
+        Returns:
+           dict[str: list]: The augmented data, where each value is wrapped
+               into a list.
+        """
+
+        aug_data = []
+        # flip_aug = [False, True] if self.flip else [False]
+        # for scale in self.img_scale:
+        #     for flip in flip_aug:
+        #         for direction in self.flip_direction:
+        #             _results = results.copy()
+        #             _results['scale'] = scale
+        #             _results['flip'] = flip
+        #             _results['flip_direction'] = direction
+        #             data = self.transforms(_results)
+        #             aug_data.append(data)
+        for _ in range(self.num):
+            _results = results.copy()
+            data = self.transforms(_results)
+            aug_data.append(data)
+        # list of dict to dict of list
+        aug_data_dict = {key: [] for key in aug_data[0]}
+        for data in aug_data:
+            for key, val in data.items():
+                aug_data_dict[key].append(val)
+        return aug_data_dict
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(transforms={self.transforms}, '
+        repr_str += f'img_scale={self.img_scale}, flip={self.flip})'
+        repr_str += f'flip_direction={self.flip_direction}'
+        return repr_str

CAGE_expression_inference-apvit/apvit_mmcls/datasets/pipelines/transforms.py

@@ -0,0 +1,918 @@
+import inspect
+import math
+import random
+import numbers
+import warnings
+from typing import Tuple, List, Optional
+
+import cv2
+import mmcv
+import numpy as np
+from PIL import Image, ImageFilter
+
+import torch
+from torch import Tensor
+from torchvision import transforms as _transforms
+import torchvision.transforms.functional as TVF
+from mmcv.utils import build_from_cfg
+
+from ..builder import PIPELINES
+
+try:
+    import albumentations
+    from albumentations import Compose
+except ImportError:
+    albumentations = None
+    Compose = None
+
+
+@PIPELINES.register_module()
+class RandomAppliedTrans(object):
+    """Randomly applied transformations.
+
+    Args:
+        transforms (list[dict]): List of transformations in dictionaries.
+        p (float): Probability.
+    """
+
+    def __init__(self, transforms, p=0.5):
+        t = [build_from_cfg(t, PIPELINES) for t in transforms]
+        self.trans = _transforms.RandomApply(t, p=p)
+
+    def __call__(self, img):
+        return self.trans(img)
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        return repr_str
+
+
+@PIPELINES.register_module()
+class RandomCrop(object):
+    """Crop the given Image at a random location.
+
+    Args:
+        size (sequence or int): Desired output size of the crop. If size is an
+            int instead of sequence like (h, w), a square crop (size, size) is
+            made.
+        padding (int or sequence, optional): Optional padding on each border
+            of the image. If a sequence of length 4 is provided, it is used to
+            pad left, top, right, bottom borders respectively.  If a sequence
+            of length 2 is provided, it is used to pad left/right, top/bottom
+            borders, respectively. Default: None, which means no padding.
+        pad_if_needed (boolean): It will pad the image if smaller than the
+            desired size to avoid raising an exception. Since cropping is done
+            after padding, the padding seems to be done at a random offset.
+            Default: False.
+        pad_val (Number | Sequence[Number]): Pixel pad_val value for constant
+            fill. If a tuple of length 3, it is used to pad_val R, G, B
+            channels respectively. Default: 0.
+        padding_mode (str): Type of padding. Should be: constant, edge,
+            reflect or symmetric. Default: constant.
+            -constant: Pads with a constant value, this value is specified
+                with pad_val.
+            -edge: pads with the last value at the edge of the image.
+            -reflect: Pads with reflection of image without repeating the
+                last value on the edge. For example, padding [1, 2, 3, 4]
+                with 2 elements on both sides in reflect mode will result
+                in [3, 2, 1, 2, 3, 4, 3, 2].
+            -symmetric: Pads with reflection of image repeating the last
+                value on the edge. For example, padding [1, 2, 3, 4] with
+                2 elements on both sides in symmetric mode will result in
+                [2, 1, 1, 2, 3, 4, 4, 3].
+    """
+
+    def __init__(self,
+                 size,
+                 padding=None,
+                 pad_if_needed=False,
+                 pad_val=0,
+                 padding_mode='constant'):
+        if isinstance(size, (tuple, list)):
+            self.size = size
+        else:
+            self.size = (size, size)
+        # check padding mode
+        assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric']
+        self.padding = padding
+        self.pad_if_needed = pad_if_needed
+        self.pad_val = pad_val
+        self.padding_mode = padding_mode
+
+    @staticmethod
+    def get_params(img, output_size):
+        """Get parameters for ``crop`` for a random crop.
+
+        Args:
+            img (ndarray): Image to be cropped.
+            output_size (tuple): Expected output size of the crop.
+
+        Returns:
+            tuple: Params (xmin, ymin, target_height, target_width) to be
+                passed to ``crop`` for random crop.
+        """
+        height = img.shape[0]
+        width = img.shape[1]
+        target_height, target_width = output_size
+        if width == target_width and height == target_height:
+            return 0, 0, height, width
+
+        xmin = random.randint(0, height - target_height)
+        ymin = random.randint(0, width - target_width)
+        return xmin, ymin, target_height, target_width
+
+    def __call__(self, results):
+        """
+        Args:
+            img (ndarray): Image to be cropped.
+        """
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            if self.padding is not None:
+                img = mmcv.impad(
+                    img, padding=self.padding, pad_val=self.pad_val)
+
+            # pad the height if needed
+            if self.pad_if_needed and img.shape[0] < self.size[0]:
+                img = mmcv.impad(
+                    img,
+                    padding=(0, self.size[0] - img.shape[0], 0,
+                             self.size[0] - img.shape[0]),
+                    pad_val=self.pad_val,
+                    padding_mode=self.padding_mode)
+
+            # pad the width if needed
+            if self.pad_if_needed and img.shape[1] < self.size[1]:
+                img = mmcv.impad(
+                    img,
+                    padding=(self.size[1] - img.shape[1], 0,
+                             self.size[1] - img.shape[1], 0),
+                    pad_val=self.pad_val,
+                    padding_mode=self.padding_mode)
+
+            xmin, ymin, height, width = self.get_params(img, self.size)
+            results[key] = mmcv.imcrop(
+                img,
+                np.array([ymin, xmin, ymin + width - 1, xmin + height - 1]))
+        return results
+
+    def __repr__(self):
+        return (self.__class__.__name__ +
+                f'(size={self.size}, padding={self.padding})')
+
+
+@PIPELINES.register_module()
+class RandomResizedCrop(object):
+    """Crop the given image to random size and aspect ratio.
+
+    A crop of random size (default: of 0.08 to 1.0) of the original size and a
+    random aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio
+    is made. This crop is finally resized to given size.
+
+    Args:
+        size (sequence or int): Desired output size of the crop. If size is an
+            int instead of sequence like (h, w), a square crop (size, size) is
+            made.
+        scale (tuple): Range of the random size of the cropped image compared
+            to the original image. Default: (0.08, 1.0).
+        ratio (tuple): Range of the random aspect ratio of the cropped image
+            compared to the original image. Default: (3. / 4., 4. / 3.).
+        interpolation (str): Interpolation method, accepted values are
+            'nearest', 'bilinear', 'bicubic', 'area', 'lanczos'. Default:
+            'bilinear'.
+        backend (str): The image resize backend type, accpeted values are
+            `cv2` and `pillow`. Default: `cv2`.
+    """
+
+    def __init__(self,
+                 size,
+                 scale=(0.08, 1.0),
+                 ratio=(3. / 4., 4. / 3.),
+                 interpolation='bilinear',
+                 backend='cv2'):
+        if isinstance(size, (tuple, list)):
+            self.size = size
+        else:
+            self.size = (size, size)
+        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
+            raise ValueError('range should be of kind (min, max). '
+                             f'But received {scale}')
+        if backend not in ['cv2', 'pillow']:
+            raise ValueError(f'backend: {backend} is not supported for resize.'
+                             'Supported backends are "cv2", "pillow"')
+
+        self.interpolation = interpolation
+        self.scale = scale
+        self.ratio = ratio
+        self.backend = backend
+
+    @staticmethod
+    def get_params(img, scale, ratio):
+        """Get parameters for ``crop`` for a random sized crop.
+
+        Args:
+            img (ndarray): Image to be cropped.
+            scale (tuple): Range of the random size of the cropped image
+                compared to the original image size.
+            ratio (tuple): Range of the random aspect ratio of the cropped
+                image compared to the original image area.
+
+        Returns:
+            tuple: Params (xmin, ymin, target_height, target_width) to be
+                passed to ``crop`` for a random sized crop.
+        """
+        height = img.shape[0]
+        width = img.shape[1]
+        area = height * width
+
+        for _ in range(10):
+            target_area = random.uniform(*scale) * area
+            log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
+            aspect_ratio = math.exp(random.uniform(*log_ratio))
+
+            target_width = int(round(math.sqrt(target_area * aspect_ratio)))
+            target_height = int(round(math.sqrt(target_area / aspect_ratio)))
+
+            if 0 < target_width <= width and 0 < target_height <= height:
+                xmin = random.randint(0, height - target_height)
+                ymin = random.randint(0, width - target_width)
+                return xmin, ymin, target_height, target_width
+
+        # Fallback to central crop
+        in_ratio = float(width) / float(height)
+        if in_ratio < min(ratio):
+            target_width = width
+            target_height = int(round(target_width / min(ratio)))
+        elif in_ratio > max(ratio):
+            target_height = height
+            target_width = int(round(target_height * max(ratio)))
+        else:  # whole image
+            target_width = width
+            target_height = height
+        xmin = (height - target_height) // 2
+        ymin = (width - target_width) // 2
+        return xmin, ymin, target_height, target_width
+
+    def __call__(self, results):
+        """
+        Args:
+            img (ndarray): Image to be cropped and resized.
+
+        Returns:
+            ndarray: Randomly cropped and resized image.
+        """
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            xmin, ymin, target_height, target_width = self.get_params(
+                img, self.scale, self.ratio)
+            img = mmcv.imcrop(
+                img,
+                np.array([
+                    ymin, xmin, ymin + target_width - 1,
+                    xmin + target_height - 1
+                ]))
+            results[key] = mmcv.imresize(
+                img,
+                tuple(self.size[::-1]),
+                interpolation=self.interpolation,
+                backend=self.backend)
+        return results
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + f'(size={self.size}'
+        format_string += f', scale={tuple(round(s, 4) for s in self.scale)}'
+        format_string += f', ratio={tuple(round(r, 4) for r in self.ratio)}'
+        format_string += f', interpolation={self.interpolation})'
+        return format_string
+
+
+@PIPELINES.register_module()
+class RandomGrayscale(object):
+    """Randomly convert image to grayscale with a probability of gray_prob.
+
+    Args:
+        gray_prob (float): Probability that image should be converted to
+            grayscale. Default: 0.1.
+
+    Returns:
+        ndarray: Grayscale version of the input image with probability
+            gray_prob and unchanged with probability (1-gray_prob).
+            - If input image is 1 channel: grayscale version is 1 channel.
+            - If input image is 3 channel: grayscale version is 3 channel
+                with r == g == b.
+
+    """
+
+    def __init__(self, gray_prob=0.1):
+        self.gray_prob = gray_prob
+
+    def __call__(self, results):
+        """
+        Args:
+            img (ndarray): Image to be converted to grayscale.
+
+        Returns:
+            ndarray: Randomly grayscaled image.
+        """
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            num_output_channels = img.shape[2]
+            if random.random() < self.gray_prob:
+                if num_output_channels > 1:
+                    img = mmcv.rgb2gray(img)[:, :, None]
+                    results[key] = np.dstack(
+                        [img for _ in range(num_output_channels)])
+                    return results
+            results[key] = img
+        return results
+
+    def __repr__(self):
+        return self.__class__.__name__ + f'(gray_prob={self.gray_prob})'
+
+
+@PIPELINES.register_module()
+class RandomFlip(object):
+    """Flip the image randomly.
+
+    Flip the image randomly based on flip probaility and flip direction.
+
+    Args:
+        flip_prob (float): probability of the image being flipped. Default: 0.5
+        direction (str, optional): The flipping direction. Options are
+            'horizontal' and 'vertical'. Default: 'horizontal'.
+    """
+
+    def __init__(self, flip_prob=0.5, direction='horizontal'):
+        assert 0 <= flip_prob <= 1
+        assert direction in ['horizontal', 'vertical']
+        self.flip_prob = flip_prob
+        self.direction = direction
+
+    def __call__(self, results):
+        """Call function to flip image.
+
+        Args:
+            results (dict): Result dict from loading pipeline.
+
+        Returns:
+            dict: Flipped results, 'flip', 'flip_direction' keys are added into
+                result dict.
+        """
+        flip = True if np.random.rand() < self.flip_prob else False
+        results['flip'] = flip
+        results['flip_direction'] = self.direction
+        if results['flip']:
+            # flip image
+            for key in results.get('img_fields', ['img']):
+                results[key] = mmcv.imflip(
+                    results[key], direction=results['flip_direction'])
+        return results
+
+    def __repr__(self):
+        return self.__class__.__name__ + f'(flip_prob={self.flip_prob})'
+
+
+@PIPELINES.register_module()
+class Resize(object):
+    """Resize images.
+
+    Args:
+        size (int | tuple): Images scales for resizing (h, w).
+            When size is int, the default behavior is to resize an image
+            to (size, size). When size is tuple and the second value is -1,
+            the short edge of an image is resized to its first value.
+            For example, when size is 224, the image is resized to 224x224.
+            When size is (224, -1), the short side is resized to 224 and the
+            other side is computed based on the short side, maintaining the
+            aspect ratio.
+        interpolation (str): Interpolation method, accepted values are
+            "nearest", "bilinear", "bicubic", "area", "lanczos".
+            More details can be found in `mmcv.image.geometric`.
+        backend (str): The image resize backend type, accpeted values are
+            `cv2` and `pillow`. Default: `cv2`.
+    """
+
+    def __init__(self, size, interpolation='bilinear', backend='cv2'):
+        assert isinstance(size, int) or (isinstance(size, tuple)
+                                         and len(size) == 2)
+        self.resize_w_short_side = False
+        if isinstance(size, int):
+            assert size > 0
+            size = (size, size)
+        else:
+            assert size[0] > 0 and (size[1] > 0 or size[1] == -1)
+            if size[1] == -1:
+                self.resize_w_short_side = True
+        assert interpolation in ('nearest', 'bilinear', 'bicubic', 'area',
+                                 'lanczos')
+        if backend not in ['cv2', 'pillow']:
+            raise ValueError(f'backend: {backend} is not supported for resize.'
+                             'Supported backends are "cv2", "pillow"')
+
+        self.size = size
+        self.interpolation = interpolation
+        self.backend = backend
+
+    def _resize_img(self, results):
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            ignore_resize = False
+            if self.resize_w_short_side:
+                h, w = img.shape[:2]
+                short_side = self.size[0]
+                if (w <= h and w == short_side) or (h <= w
+                                                    and h == short_side):
+                    ignore_resize = True
+                else:
+                    if w < h:
+                        width = short_side
+                        height = int(short_side * h / w)
+                    else:
+                        height = short_side
+                        width = int(short_side * w / h)
+            else:
+                height, width = self.size
+            if not ignore_resize:
+                img = mmcv.imresize(
+                    img,
+                    size=(width, height),
+                    interpolation=self.interpolation,
+                    return_scale=False,
+                    backend=self.backend)
+                results[key] = img
+                results['img_shape'] = img.shape
+
+    def __call__(self, results):
+        self._resize_img(results)
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(size={self.size}, '
+        repr_str += f'interpolation={self.interpolation})'
+        return repr_str
+
+
+@PIPELINES.register_module()
+class CenterCrop(object):
+    """Center crop the image.
+
+    Args:
+        crop_size (int | tuple): Expected size after cropping, (h, w).
+
+    Notes:
+        If the image is smaller than the crop size, return the original image
+    """
+
+    def __init__(self, crop_size):
+        assert isinstance(crop_size, int) or (isinstance(crop_size, tuple)
+                                              and len(crop_size) == 2)
+        if isinstance(crop_size, int):
+            crop_size = (crop_size, crop_size)
+        assert crop_size[0] > 0 and crop_size[1] > 0
+        self.crop_size = crop_size
+
+    def __call__(self, results):
+        crop_height, crop_width = self.crop_size[0], self.crop_size[1]
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            # img.shape has length 2 for grayscale, length 3 for color
+            img_height, img_width = img.shape[:2]
+
+            y1 = max(0, int(round((img_height - crop_height) / 2.)))
+            x1 = max(0, int(round((img_width - crop_width) / 2.)))
+            y2 = min(img_height, y1 + crop_height) - 1
+            x2 = min(img_width, x1 + crop_width) - 1
+
+            # crop the image
+            img = mmcv.imcrop(img, bboxes=np.array([x1, y1, x2, y2]))
+            img_shape = img.shape
+            results[key] = img
+        results['img_shape'] = img_shape
+
+        return results
+
+    def __repr__(self):
+        return self.__class__.__name__ + f'(crop_size={self.crop_size})'
+
+
+@PIPELINES.register_module()
+class Normalize(object):
+    """Normalize the image.
+
+    Args:
+        mean (sequence): Mean values of 3 channels.
+        std (sequence): Std values of 3 channels.
+        to_rgb (bool): Whether to convert the image from BGR to RGB,
+            default is true.
+    """
+
+    def __init__(self, mean, std, to_rgb=True):
+        self.mean = np.array(mean, dtype=np.float32)
+        self.std = np.array(std, dtype=np.float32)
+        self.to_rgb = to_rgb
+
+    def __call__(self, results):
+        for key in results.get('img_fields', ['img']):
+            results[key] = mmcv.imnormalize(results[key], self.mean, self.std,
+                                            self.to_rgb)
+        results['img_norm_cfg'] = dict(
+            mean=self.mean, std=self.std, to_rgb=self.to_rgb)
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(mean={list(self.mean)}, '
+        repr_str += f'std={list(self.std)}, '
+        repr_str += f'to_rgb={self.to_rgb})'
+        return repr_str
+
+
+@PIPELINES.register_module()
+class Albu(object):
+    """Albumentation augmentation.
+
+    Adds custom transformations from Albumentations library.
+    Please, visit `https://albumentations.readthedocs.io`
+    to get more information.
+    An example of ``transforms`` is as followed:
+
+    .. code-block::
+        [
+            dict(
+                type='ShiftScaleRotate',
+                shift_limit=0.0625,
+                scale_limit=0.0,
+                rotate_limit=0,
+                interpolation=1,
+                p=0.5),
+            dict(
+                type='RandomBrightnessContrast',
+                brightness_limit=[0.1, 0.3],
+                contrast_limit=[0.1, 0.3],
+                p=0.2),
+            dict(type='ChannelShuffle', p=0.1),
+            dict(
+                type='OneOf',
+                transforms=[
+                    dict(type='Blur', blur_limit=3, p=1.0),
+                    dict(type='MedianBlur', blur_limit=3, p=1.0)
+                ],
+                p=0.1),
+        ]
+
+    Args:
+        transforms (list[dict]): A list of albu transformations
+        keymap (dict): Contains {'input key':'albumentation-style key'}
+    """
+
+    def __init__(self, transforms, keymap=None, update_pad_shape=False):
+        if Compose is None:
+            raise RuntimeError('albumentations is not installed')
+
+        self.transforms = transforms
+        self.filter_lost_elements = False
+        self.update_pad_shape = update_pad_shape
+
+        self.aug = Compose([self.albu_builder(t) for t in self.transforms])
+
+        if not keymap:
+            self.keymap_to_albu = {
+                'img': 'image',
+            }
+        else:
+            self.keymap_to_albu = keymap
+        self.keymap_back = {v: k for k, v in self.keymap_to_albu.items()}
+
+    def albu_builder(self, cfg):
+        """Import a module from albumentations.
+        It inherits some of :func:`build_from_cfg` logic.
+        Args:
+            cfg (dict): Config dict. It should at least contain the key "type".
+        Returns:
+            obj: The constructed object.
+        """
+
+        assert isinstance(cfg, dict) and 'type' in cfg
+        args = cfg.copy()
+
+        obj_type = args.pop('type')
+        if mmcv.is_str(obj_type):
+            if albumentations is None:
+                raise RuntimeError('albumentations is not installed')
+            obj_cls = getattr(albumentations, obj_type)
+        elif inspect.isclass(obj_type):
+            obj_cls = obj_type
+        else:
+            raise TypeError(
+                f'type must be a str or valid type, but got {type(obj_type)}')
+
+        if 'transforms' in args:
+            args['transforms'] = [
+                self.albu_builder(transform)
+                for transform in args['transforms']
+            ]
+
+        return obj_cls(**args)
+
+    @staticmethod
+    def mapper(d, keymap):
+        """Dictionary mapper. Renames keys according to keymap provided.
+        Args:
+            d (dict): old dict
+            keymap (dict): {'old_key':'new_key'}
+        Returns:
+            dict: new dict.
+        """
+
+        updated_dict = {}
+        for k, v in zip(d.keys(), d.values()):
+            new_k = keymap.get(k, k)
+            updated_dict[new_k] = d[k]
+        return updated_dict
+
+    def __call__(self, results):
+        # dict to albumentations format
+        results = self.mapper(results, self.keymap_to_albu)
+
+        results = self.aug(**results)
+
+        if 'gt_labels' in results:
+            if isinstance(results['gt_labels'], list):
+                results['gt_labels'] = np.array(results['gt_labels'])
+            results['gt_labels'] = results['gt_labels'].astype(np.int64)
+
+        # back to the original format
+        results = self.mapper(results, self.keymap_back)
+
+        # update final shape
+        if self.update_pad_shape:
+            results['pad_shape'] = results['img'].shape
+
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__ + f'(transforms={self.transforms})'
+        return repr_str
+
+# custom transforms
+@PIPELINES.register_module()
+class GaussianBlur(object):
+
+    def __init__(self, sigma_min, sigma_max):
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+
+    def __call__(self, results):
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            sigma = np.random.uniform(self.sigma_min, self.sigma_max)
+            img = Image.fromarray(img)
+            img = img.filter(ImageFilter.GaussianBlur(radius=sigma))
+            results[key] = np.array(img).astype('float32')
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        return repr_str
+
+
+@PIPELINES.register_module()
+class RandomRotate(object):
+    """Rotate the image & seg.
+
+    Args:
+        prob (float): The rotation probability.
+        degree (float, tuple[float]): Range of degrees to select from. If
+            degree is a number instead of tuple like (min, max),
+            the range of degree will be (``-degree``, ``+degree``)
+        pad_val (float, optional): Padding value of image. Default: 0.
+        center (tuple[float], optional): Center point (w, h) of the rotation in
+            the source image. If not specified, the center of the image will be
+            used. Default: None.
+        auto_bound (bool): Whether to adjust the image size to cover the whole
+            rotated image. Default: False
+    """
+
+    def __init__(self,
+                 prob,
+                 degree,
+                 pad_val=0,
+                 center=None,
+                 auto_bound=False):
+        self.prob = prob
+        assert prob >= 0 and prob <= 1
+        if isinstance(degree, (float, int)):
+            assert degree > 0, f'degree {degree} should be positive'
+            self.degree = (-degree, degree)
+        else:
+            self.degree = degree
+        assert len(self.degree) == 2, f'degree {self.degree} should be a ' \
+                                      f'tuple of (min, max)'
+        self.pal_val = pad_val
+        self.center = center
+        self.auto_bound = auto_bound
+
+    def __call__(self, results):
+        """Call function to rotate image.
+
+        Args:
+            results (dict): Result dict from loading pipeline.
+
+        Returns:
+            dict: Rotated results.
+        """
+
+        rotate = True if np.random.rand() < self.prob else False
+        degree = np.random.uniform(min(*self.degree), max(*self.degree))
+        if rotate:
+            # rotate image
+            for key in results.get('img_fields', ['img']):
+                results[key] = mmcv.imrotate(
+                    results[key],
+                    angle=degree,
+                    border_value=self.pal_val,
+                    center=self.center,
+                    auto_bound=self.auto_bound)
+
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(prob={self.prob}, ' \
+                    f'degree={self.degree}, ' \
+                    f'pad_val={self.pal_val}, ' \
+                    f'center={self.center}, ' \
+                    f'auto_bound={self.auto_bound})'
+        return repr_str
+
+
+@PIPELINES.register_module()
+class RandomErasing:
+    """ Randomly selects a rectangle region in an image and erases its pixels.
+    'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896
+
+    Args:
+         p: probability that the random erasing operation will be performed.
+         scale: range of proportion of erased area against input image.
+         ratio: range of aspect ratio of erased area.
+         value: erasing value. Default is 0. If a single int, it is used to
+            erase all pixels. If a tuple of length 3, it is used to erase
+            R, G, B channels respectively.
+            If a str of 'random', erasing each pixel with random values.
+         inplace: boolean to make this transform inplace. Default set to False.
+
+    Returns:
+        Erased Image.
+    """
+    def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False):
+        super().__init__()
+        if not isinstance(value, (numbers.Number, str, tuple, list)):
+            raise TypeError("Argument value should be either a number or str or a sequence")
+        if isinstance(value, str) and value != "random":
+            raise ValueError("If value is str, it should be 'random'")
+        if not isinstance(scale, (tuple, list)):
+            raise TypeError("Scale should be a sequence")
+        if not isinstance(ratio, (tuple, list)):
+            raise TypeError("Ratio should be a sequence")
+        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
+            warnings.warn("Scale and ratio should be of kind (min, max)")
+        if scale[0] < 0 or scale[1] > 1:
+            raise ValueError("Scale should be between 0 and 1")
+        if p < 0 or p > 1:
+            raise ValueError("Random erasing probability should be between 0 and 1")
+
+        self.p = p
+        self.scale = scale
+        self.ratio = ratio
+        self.value = value
+        self.inplace = inplace
+
+    @staticmethod
+    def get_params(
+            img: Tensor, scale: Tuple[float, float], ratio: Tuple[float, float], value: Optional[List[float]] = None
+        ) -> Tuple[int, int, int, int, Tensor]:
+        """Get parameters for ``erase`` for a random erasing.
+
+        Args:
+            img (Tensor): Tensor image to be erased.
+            scale (tuple or list): range of proportion of erased area against input image.
+            ratio (tuple or list): range of aspect ratio of erased area.
+            value (list, optional): erasing value. If None, it is interpreted as "random"
+                (erasing each pixel with random values). If ``len(value)`` is 1, it is interpreted as a number,
+                i.e. ``value[0]``.
+
+        Returns:
+            tuple: params (i, j, h, w, v) to be passed to ``erase`` for random erasing.
+        """
+        img_c, img_h, img_w = img.shape[-3], img.shape[-2], img.shape[-1]
+        area = img_h * img_w
+
+        for _ in range(10):
+            erase_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
+            aspect_ratio = torch.empty(1).uniform_(ratio[0], ratio[1]).item()
+
+            h = int(round(math.sqrt(erase_area * aspect_ratio)))
+            w = int(round(math.sqrt(erase_area / aspect_ratio)))
+            if not (h < img_h and w < img_w):
+                continue
+
+            if value is None:
+                v = torch.empty([img_c, h, w], dtype=torch.float32).normal_()
+            else:
+                v = torch.tensor(value)[:, None, None]
+
+            i = torch.randint(0, img_h - h + 1, size=(1, )).item()
+            j = torch.randint(0, img_w - w + 1, size=(1, )).item()
+            return i, j, h, w, v
+
+        # Return original image
+        return 0, 0, img_h, img_w, img
+    
+    def __call__(self, results):
+        if torch.rand(1) < self.p:
+            for key in results.get('img_fields', ['img']):
+                img = results[key]
+                img = torch.Tensor(img.transpose([2,0,1]))
+                # the next process need C, H, W
+
+                # cast self.value to script acceptable type
+                if isinstance(self.value, (int, float)):
+                    value = [self.value, ]
+                elif isinstance(self.value, str):
+                    value = None
+                elif isinstance(self.value, tuple):
+                    value = list(self.value)
+                else:
+                    value = self.value
+
+                if value is not None and not (len(value) in (1, img.shape[-3])):
+                    raise ValueError(
+                        "If value is a sequence, it should have either a single value or "
+                        "{} (number of input channels)".format(img.shape[-3])
+                    )
+
+                x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=value)
+                results[key] = TVF.erase(img, x, y, h, w, v, self.inplace).numpy().transpose([1,2,0])
+        return results
+
+
+@PIPELINES.register_module()
+class ColorJitter:
+    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
+        self.func = _transforms.ColorJitter(brightness, contrast, saturation, hue)
+    
+    def __call__(self, results):
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            img = Image.fromarray(img)
+            results[key] = np.array(self.func(img))
+        return results
+    
+    def __repr__(self):
+        return self.func.__repr__()
+
+
+@PIPELINES.register_module()
+class HistogramEqualization:
+
+    def __call__(self, results):
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            # convert from RGB color-space to YCrCb
+            ycrcb_img = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
+
+            # equalize the histogram of the Y channel
+            ycrcb_img[:, :, 0] = cv2.equalizeHist(ycrcb_img[:, :, 0])
+
+            # convert back to RGB color-space from YCrCb
+            equalized_img = cv2.cvtColor(ycrcb_img, cv2.COLOR_YCrCb2BGR)
+            results[key] = equalized_img
+        return results
+
+
+@PIPELINES.register_module()
+class TorchNormalize:
+    """
+    Normalize in torchvision
+    """
+    def __init__(self, mean, std, inplace=False):
+        self.mean = mean
+        self.std = std
+        self.inplace = inplace
+
+    def __call__(self, results):
+        """
+        Args:
+            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
+
+        Returns:
+            Tensor: Normalized Tensor image.
+        """
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            results[key] = _transforms.functional.normalize(img, self.mean, self.std, self.inplace)
+        return results
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)

CAGE_expression_inference-apvit/apvit_mmcls/datasets/raf.py

@@ -0,0 +1,144 @@
+import os
+import random
+from typing import Any, Dict
+
+import numpy as np
+from mmcls.models.losses import accuracy, f1_score, precision, recall
+from mmcls.models.losses.eval_metrics import class_accuracy
+
+from .base_dataset import BaseDataset
+from .builder import DATASETS
+
+
+def has_file_allowed_extension(filename, extensions):
+    """Checks if a file is an allowed extension.
+
+    Args:
+        filename (string): path to a file
+
+    Returns:
+        bool: True if the filename ends with a known image extension
+    """
+    filename_lower = filename.lower()
+    return any(filename_lower.endswith(ext) for ext in extensions)
+
+
+def find_folders(root):
+    """Find classes by folders under a root.
+
+    Args:
+        root (string): root directory of folders
+
+    Returns:
+        folder_to_idx (dict): the map from folder name to class idx
+    """
+    folders = [
+        d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))
+    ]
+    folders.sort()
+    folder_to_idx = {folders[i]: i for i in range(len(folders))}
+    return folder_to_idx
+
+
+def get_samples(root, folder_to_idx, extensions):
+    """Make dataset by walking all images under a root.
+
+    Args:
+        root (string): root directory of folders
+        folder_to_idx (dict): the map from class name to class idx
+        extensions (tuple): allowed extensions
+
+    Returns:
+        samples (list): a list of tuple where each element is (image, label)
+    """
+    samples = []
+    root = os.path.expanduser(root)
+    for folder_name in sorted(os.listdir(root)):
+        _dir = os.path.join(root, folder_name)
+        if not os.path.isdir(_dir):
+            continue
+
+        for _, _, fns in sorted(os.walk(_dir)):
+            for fn in sorted(fns):
+                if has_file_allowed_extension(fn, extensions):
+                    path = os.path.join(folder_name, fn)
+                    item = (path, folder_to_idx[folder_name])
+                    samples.append(item)
+    return samples
+
+#                  0       1         2          3       4               5           6           7
+FER_CLASSES = ['Anger', 'Disgust', 'Fear', 'Sadness', 'Happiness', 'Surprise', 'Neutral', 'Contempt']
+
+def convert2coarse_label(i:int):
+    """The first four are negative"""
+    if i <= 3:
+        return 0
+    return i - 3
+
+
+def gen_class_map(dataset_class):
+    """
+    generate the convert map from DATASET_CLASSES to FER_CLASSES
+    """
+    convert_map = []
+    for i in dataset_class:
+        convert_map.append(FER_CLASSES.index(i))
+    assert sum(convert_map) == sum([i for i in range(len(dataset_class))])
+    return convert_map
+
+@DATASETS.register_module()
+class RAF(BaseDataset):
+
+    IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif')
+    DATASET_CLASSES = [
+        'Surprise',
+        'Fear',
+        'Disgust',
+        'Happiness',
+        'Sadness',
+        'Anger',
+        'Neutral'
+    ]
+    CLASSES = FER_CLASSES[:7]
+
+    @staticmethod
+    def convert_gt_label(i:int):
+        """# dataset -> FER_CLASSES"""
+        convert_table = (5, 2, 1, 4, 3, 0, 6)
+        assert sum(convert_table) == sum([i for i in range(7)])
+        return convert_table[i]
+
+    def load_annotations(self):
+        if self.ann_file is None:
+            folder_to_idx = find_folders(self.data_prefix)
+            samples = get_samples(
+                self.data_prefix,
+                folder_to_idx,
+                extensions=self.IMG_EXTENSIONS)
+            if len(samples) == 0:
+                raise (RuntimeError('Found 0 files in subfolders of: '
+                                    f'{self.data_prefix}. '
+                                    'Supported extensions are: '
+                                    f'{",".join(self.IMG_EXTENSIONS)}'))
+
+            self.folder_to_idx = folder_to_idx
+        elif isinstance(self.ann_file, str):
+            with open(self.ann_file) as f:
+                samples = [x.strip().split(' ') for x in f.readlines()]
+            samples = [[i[0].replace('_aligned', ''), i[1]] for i in samples]
+        else:
+            raise TypeError('ann_file must be a str or None')
+        self.samples = samples
+
+        data_infos = []
+        for filename, gt_label in self.samples:
+            info = {'img_prefix': self.data_prefix}
+            info['img_info'] = {'filename': filename}
+            gt_label = int(gt_label) - 1
+            gt_label = self.convert_gt_label(gt_label)
+            coarse_label = convert2coarse_label(gt_label)
+            info['gt_label'] = np.array(gt_label, dtype=np.int64)
+            info['coarse_label'] = np.array(coarse_label, dtype=np.int64)
+            data_infos.append(info)
+        return data_infos
+

CAGE_expression_inference-apvit/apvit_mmcls/datasets/samplers/__init__.py

@@ -0,0 +1,3 @@
+from .distributed_sampler import DistributedSampler
+
+__all__ = ['DistributedSampler']

CAGE_expression_inference-apvit/apvit_mmcls/datasets/samplers/distributed_sampler.py

@@ -0,0 +1,42 @@
+import torch
+from torch.utils.data import DistributedSampler as _DistributedSampler
+
+
+class DistributedSampler(_DistributedSampler):
+
+    def __init__(self,
+                 dataset,
+                 num_replicas=None,
+                 rank=None,
+                 shuffle=True,
+                 round_up=True):
+        super().__init__(dataset, num_replicas=num_replicas, rank=rank)
+        self.shuffle = shuffle
+        self.round_up = round_up
+        if self.round_up:
+            self.total_size = self.num_samples * self.num_replicas
+        else:
+            self.total_size = len(self.dataset)
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = torch.arange(len(self.dataset)).tolist()
+
+        # add extra samples to make it evenly divisible
+        if self.round_up:
+            indices = (
+                indices *
+                int(self.total_size / len(indices) + 1))[:self.total_size]
+        assert len(indices) == self.total_size
+
+        # subsample
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        if self.round_up:
+            assert len(indices) == self.num_samples
+
+        return iter(indices)

CAGE_expression_inference-apvit/apvit_mmcls/datasets/utils.py

@@ -0,0 +1,152 @@
+import gzip
+import hashlib
+import os
+import os.path
+import shutil
+import tarfile
+import urllib.error
+import urllib.request
+import zipfile
+
+__all__ = ['rm_suffix', 'check_integrity', 'download_and_extract_archive']
+
+
+def rm_suffix(s, suffix=None):
+    if suffix is None:
+        return s[:s.rfind('.')]
+    else:
+        return s[:s.rfind(suffix)]
+
+
+def calculate_md5(fpath, chunk_size=1024 * 1024):
+    md5 = hashlib.md5()
+    with open(fpath, 'rb') as f:
+        for chunk in iter(lambda: f.read(chunk_size), b''):
+            md5.update(chunk)
+    return md5.hexdigest()
+
+
+def check_md5(fpath, md5, **kwargs):
+    return md5 == calculate_md5(fpath, **kwargs)
+
+
+def check_integrity(fpath, md5=None):
+    if not os.path.isfile(fpath):
+        return False
+    if md5 is None:
+        return True
+    return check_md5(fpath, md5)
+
+
+def download_url_to_file(url, fpath):
+    with urllib.request.urlopen(url) as resp, open(fpath, 'wb') as of:
+        shutil.copyfileobj(resp, of)
+
+
+def download_url(url, root, filename=None, md5=None):
+    """Download a file from a url and place it in root.
+
+    Args:
+        url (str): URL to download file from.
+        root (str): Directory to place downloaded file in.
+        filename (str | None): Name to save the file under.
+            If filename is None, use the basename of the URL.
+        md5 (str | None): MD5 checksum of the download.
+            If md5 is None, download without md5 check.
+    """
+    root = os.path.expanduser(root)
+    if not filename:
+        filename = os.path.basename(url)
+    fpath = os.path.join(root, filename)
+
+    os.makedirs(root, exist_ok=True)
+
+    if check_integrity(fpath, md5):
+        print(f'Using downloaded and verified file: {fpath}')
+    else:
+        try:
+            print(f'Downloading {url} to {fpath}')
+            download_url_to_file(url, fpath)
+        except (urllib.error.URLError, IOError) as e:
+            if url[:5] == 'https':
+                url = url.replace('https:', 'http:')
+                print('Failed download. Trying https -> http instead.'
+                      f' Downloading {url} to {fpath}')
+                download_url_to_file(url, fpath)
+            else:
+                raise e
+        # check integrity of downloaded file
+        if not check_integrity(fpath, md5):
+            raise RuntimeError('File not found or corrupted.')
+
+
+def _is_tarxz(filename):
+    return filename.endswith('.tar.xz')
+
+
+def _is_tar(filename):
+    return filename.endswith('.tar')
+
+
+def _is_targz(filename):
+    return filename.endswith('.tar.gz')
+
+
+def _is_tgz(filename):
+    return filename.endswith('.tgz')
+
+
+def _is_gzip(filename):
+    return filename.endswith('.gz') and not filename.endswith('.tar.gz')
+
+
+def _is_zip(filename):
+    return filename.endswith('.zip')
+
+
+def extract_archive(from_path, to_path=None, remove_finished=False):
+    if to_path is None:
+        to_path = os.path.dirname(from_path)
+
+    if _is_tar(from_path):
+        with tarfile.open(from_path, 'r') as tar:
+            tar.extractall(path=to_path)
+    elif _is_targz(from_path) or _is_tgz(from_path):
+        with tarfile.open(from_path, 'r:gz') as tar:
+            tar.extractall(path=to_path)
+    elif _is_tarxz(from_path):
+        with tarfile.open(from_path, 'r:xz') as tar:
+            tar.extractall(path=to_path)
+    elif _is_gzip(from_path):
+        to_path = os.path.join(
+            to_path,
+            os.path.splitext(os.path.basename(from_path))[0])
+        with open(to_path, 'wb') as out_f, gzip.GzipFile(from_path) as zip_f:
+            out_f.write(zip_f.read())
+    elif _is_zip(from_path):
+        with zipfile.ZipFile(from_path, 'r') as z:
+            z.extractall(to_path)
+    else:
+        raise ValueError(f'Extraction of {from_path} not supported')
+
+    if remove_finished:
+        os.remove(from_path)
+
+
+def download_and_extract_archive(url,
+                                 download_root,
+                                 extract_root=None,
+                                 filename=None,
+                                 md5=None,
+                                 remove_finished=False):
+    download_root = os.path.expanduser(download_root)
+    if extract_root is None:
+        extract_root = download_root
+    if not filename:
+        filename = os.path.basename(url)
+
+    download_url(url, download_root, filename, md5)
+
+    archive = os.path.join(download_root, filename)
+    print(f'Extracting {archive} to {extract_root}')
+    extract_archive(archive, extract_root, remove_finished)

CAGE_expression_inference-apvit/apvit_mmcls/models/__init__.py

@@ -0,0 +1,13 @@
+from .backbones import *  # noqa: F401,F403
+from .builder import (BACKBONES, CLASSIFIERS, HEADS, LOSSES, NECKS,
+                      build_backbone, build_classifier, build_head, build_loss,
+                      build_neck)
+from .classifiers import *  # noqa: F401,F403
+from .heads import *  # noqa: F401,F403
+from .losses import *  # noqa: F401,F403
+from .necks import *  # noqa: F401,F403
+
+__all__ = [
+    'BACKBONES', 'HEADS', 'NECKS', 'LOSSES', 'CLASSIFIERS', 'build_backbone',
+    'build_head', 'build_neck', 'build_loss', 'build_classifier'
+]

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/__init__.py

@@ -0,0 +1,32 @@
+from .alexnet import AlexNet
+from .lenet import LeNet5
+from .mobilenet_v2 import MobileNetV2
+from .mobilenet_v3 import MobileNetv3
+from .regnet import RegNet
+from .resnest import ResNeSt
+from .resnet import ResNet, ResNetV1d
+from .resnet_cifar import ResNet_CIFAR
+from .resnext import ResNeXt
+from .seresnet import SEResNet
+from .seresnext import SEResNeXt
+from .shufflenet_v1 import ShuffleNetV1
+from .shufflenet_v2 import ShuffleNetV2
+from .vgg import VGG
+from .irse import IRSE
+from .irse_nopadding import IRSENoPadding
+from .mobilefacenet import MobileFaceNet
+from ..vit.vit_origin import VisionTransformerOrigin
+from ..vit.vit_siam_merge import PoolingViT
+from .t2t_vit import T2T_ViT, T2T_ViTPooling
+from .iresnet import IResNet
+from .vtff import VTFF, MViT
+
+__all__ = [
+    'VTFF', 'MViT',
+    'LeNet5', 'AlexNet', 'VGG', 'RegNet', 'ResNet', 'ResNeXt', 'ResNetV1d',
+    'ResNeSt', 'ResNet_CIFAR', 'SEResNet', 'SEResNeXt', 'ShuffleNetV1',
+    'ShuffleNetV2', 'MobileNetV2', 'MobileNetv3', 'IRSE',
+    'MobileFaceNet',
+    'T2T_ViT', 'T2T_ViTPooling', 'VisionTransformerOrigin', 'PoolingViT',
+    'IRSENoPadding', 'IResNet',
+]

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/alexnet.py

@@ -0,0 +1,55 @@
+import torch.nn as nn
+
+from ..builder import BACKBONES
+from .base_backbone import BaseBackbone
+
+
+@BACKBONES.register_module()
+class AlexNet(BaseBackbone):
+    """`AlexNet <https://en.wikipedia.org/wiki/AlexNet>`_ backbone.
+
+    The input for AlexNet is a 224x224 RGB image.
+
+    Args:
+        num_classes (int): number of classes for classification.
+            The default value is -1, which uses the backbone as
+            a feature extractor without the top classifier.
+    """
+
+    def __init__(self, num_classes=-1):
+        super(AlexNet, self).__init__()
+        self.num_classes = num_classes
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            nn.Conv2d(64, 192, kernel_size=5, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            nn.Conv2d(192, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+        )
+        if self.num_classes > 0:
+            self.classifier = nn.Sequential(
+                nn.Dropout(),
+                nn.Linear(256 * 6 * 6, 4096),
+                nn.ReLU(inplace=True),
+                nn.Dropout(),
+                nn.Linear(4096, 4096),
+                nn.ReLU(inplace=True),
+                nn.Linear(4096, num_classes),
+            )
+
+    def forward(self, x):
+
+        x = self.features(x)
+        if self.num_classes > 0:
+            x = x.view(x.size(0), 256 * 6 * 6)
+            x = self.classifier(x)
+
+        return x

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/base_backbone.py

@@ -0,0 +1,57 @@
+from abc import ABCMeta, abstractmethod
+
+import torch.nn as nn
+from mmcv.runner import load_checkpoint
+from mmcv.utils import get_logger
+# from mmcls.utils import get_root_logger
+
+
+class BaseBackbone(nn.Module, metaclass=ABCMeta):
+    """Base backbone.
+
+    This class defines the basic functions of a backbone.
+    Any backbone that inherits this class should at least
+    define its own `forward` function.
+
+    """
+
+    def __init__(self):
+        super(BaseBackbone, self).__init__()
+
+    def init_weights(self, pretrained=None):
+        """Init backbone weights
+
+        Args:
+            pretrained (str | None): If pretrained is a string, then it
+                initializes backbone weights by loading the pretrained
+                checkpoint. If pretrained is None, then it follows default
+                initializer or customized initializer in subclasses.
+        """
+        if isinstance(pretrained, str):
+            logger = get_logger('mmcv')
+            logger.warning(f'{self.__class__.__name__} load pretrain from {pretrained}')
+            load_checkpoint(self, pretrained, strict=False, logger=logger, map_location='cpu')
+        elif pretrained is None:
+            # use default initializer or customized initializer in subclasses
+            pass
+        else:
+            raise TypeError('pretrained must be a str or None.'
+                            f' But received {type(pretrained)}.')
+
+    @abstractmethod
+    def forward(self, x):
+        """Forward computation
+
+        Args:
+            x (tensor | tuple[tensor]): x could be a Torch.tensor or a tuple of
+                Torch.tensor, containing input data for forward computation.
+        """
+        pass
+
+    def train(self, mode=True):
+        """Set module status before forward computation
+
+        Args:
+            mode (bool): Whether it is train_mode or test_mode
+        """
+        super(BaseBackbone, self).train(mode)

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/iresnet.py

@@ -0,0 +1,220 @@
+import warnings
+import torch.nn as nn
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, ReLU, Sigmoid, Dropout, MaxPool2d, \
+    AdaptiveAvgPool2d, Sequential, Module
+from collections import namedtuple
+import torch
+from mmcv.runner import load_state_dict
+from mmcls.utils import get_root_logger
+
+from ..builder import BACKBONES
+from .base_backbone import BaseBackbone
+
+__all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200']
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=3,
+                     stride=stride,
+                     padding=dilation,
+                     groups=groups,
+                     bias=False,
+                     dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=1,
+                     stride=stride,
+                     bias=False)
+
+
+class IBasicBlock(nn.Module):
+    expansion = 1
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 groups=1, base_width=64, dilation=1):
+        super(IBasicBlock, self).__init__()
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,)
+        self.conv1 = conv3x3(inplanes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,)
+        self.prelu = nn.PReLU(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+        out = self.bn1(x)
+        out = self.conv1(out)
+        out = self.bn2(out)
+        out = self.prelu(out)
+        out = self.conv2(out)
+        out = self.bn3(out)
+        if self.downsample is not None:
+            identity = self.downsample(x)
+        out += identity
+        return out
+
+
+layers_map = {
+    18: [2, 2, 2, 2],
+    34: [3, 4, 6, 3],
+    50: [3, 4, 14, 3],
+    100: [3, 13, 30, 3],
+}
+
+
+@BACKBONES.register_module()
+class IResNet(nn.Module):
+    fc_scale = 7 * 7
+    def __init__(self,
+                 depth, block=None, dropout=0, num_features=512, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False, pretrained=None):
+        super(IResNet, self).__init__()
+        block = IBasicBlock
+        layers = layers_map[depth]
+        self.fp16 = fp16
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
+        self.prelu = nn.PReLU(self.inplanes)
+        self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
+        self.layer2 = self._make_layer(block,
+                                       128,
+                                       layers[1],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block,
+                                       256,
+                                       layers[2],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        # self.layer4 = self._make_layer(block,
+        #                                512,
+        #                                layers[3],
+        #                                stride=2,
+        #                                dilate=replace_stride_with_dilation[2])
+        # self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,)
+        # self.dropout = nn.Dropout(p=dropout, inplace=True)
+        # self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)
+        # self.features = nn.BatchNorm1d(num_features, eps=1e-05)
+        # nn.init.constant_(self.features.weight, 1.0)
+        # self.features.weight.requires_grad = False
+        if pretrained:
+            self.init_weights(pretrained)
+        else:
+            for m in self.modules():
+                if isinstance(m, nn.Conv2d):
+                    nn.init.normal_(m.weight, 0, 0.1)
+                elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                    nn.init.constant_(m.weight, 1)
+                    nn.init.constant_(m.bias, 0)
+
+            if zero_init_residual:
+                for m in self.modules():
+                    if isinstance(m, IBasicBlock):
+                        nn.init.constant_(m.bn2.weight, 0)
+
+    def init_weights(self, pretrained):
+        logger = get_root_logger()
+        logger.warning(f'{self.__class__.__name__} load pretrain from {pretrained}')
+        state_dict = torch.load(pretrained, map_location='cpu')
+        if 'state_dict' in state_dict:
+            state_dict = state_dict['state_dict']
+
+        load_state_dict(self, state_dict, strict=False, logger=logger)
+
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),
+            )
+        layers = []
+        layers.append(
+            block(self.inplanes, planes, stride, downsample, self.groups,
+                  self.base_width, previous_dilation))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(self.inplanes,
+                      planes,
+                      groups=self.groups,
+                      base_width=self.base_width,
+                      dilation=self.dilation))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        # with torch.cuda.amp.autocast(self.fp16):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.prelu(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        return (x, )
+        x = self.layer4(x)
+        return (x, )
+        #     x = self.bn2(x)
+        #     x = torch.flatten(x, 1)
+        #     x = self.dropout(x)
+        # x = self.fc(x.float() if self.fp16 else x)
+        # x = self.features(x)
+        # return x
+
+
+def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
+    model = IResNet(block, layers, **kwargs)
+    if pretrained:
+        raise ValueError()
+    return model
+
+
+def iresnet18(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet34(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet50(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet100(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet200(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained,
+                    progress, **kwargs)

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/irse.py

@@ -0,0 +1,387 @@
+import warnings
+import torch.nn as nn
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, ReLU, Sigmoid, Dropout, MaxPool2d, \
+    AdaptiveAvgPool2d, Sequential, Module
+from collections import namedtuple
+import torch
+from mmcv.runner import load_state_dict
+from mmcls.utils import get_root_logger
+
+from ..builder import BACKBONES
+from .base_backbone import BaseBackbone
+
+# Support: ['IR_50', 'IR_101', 'IR_152', 'IR_SE_50', 'IR_SE_101', 'IR_SE_152']
+
+class ABSPool2d(Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.args = args
+        self.kwargs = kwargs
+
+    def forward(self, x):
+        positive = MaxPool2d(*self.args, **self.kwargs)(x)
+        negtive = MaxPool2d(*self.args, **self.kwargs)(-x)
+        result = torch.where(positive>=negtive, positive, negtive)
+        return result
+
+pool_layer = MaxPool2d
+
+class Flatten(Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+
+def l2_norm(input, axis=1):
+    norm = torch.norm(input, 2, axis, True)
+    output = torch.div(input, norm)
+
+    return output
+
+
+class SEModule(Module):
+    def __init__(self, channels, reduction):
+        super(SEModule, self).__init__()
+        self.avg_pool = AdaptiveAvgPool2d(1)
+        self.fc1 = Conv2d(
+            channels, channels // reduction, kernel_size=1, padding=0, bias=False)
+
+        nn.init.xavier_uniform_(self.fc1.weight.data)
+
+        self.relu = ReLU(inplace=True)
+        self.fc2 = Conv2d(
+            channels // reduction, channels, kernel_size=1, padding=0, bias=False)
+
+        self.sigmoid = Sigmoid()
+
+    def forward(self, x):
+        module_input = x
+        x = self.avg_pool(x)
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        x = self.sigmoid(x)
+
+        return module_input * x
+
+
+class bottleneck_IR(Module):
+    def __init__(self, in_channel, depth, stride):
+        super(bottleneck_IR, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = pool_layer(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+                BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
+            PReLU(depth),
+            Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
+            BatchNorm2d(depth))
+
+    def forward(self, x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+
+        return res + shortcut
+
+
+class BasicBlockIR(Module):
+    """
+    BasicBlock for IRNet, stolen from TFace
+    https://github.com/Tencent/TFace/blob/d57fd8d9ce9502240921f0998c57f84afa7eaeaa/torchkit/backbone/model_irse.py#L16
+    Add a BatchNorm2d after the first Conv2d
+    """
+    def __init__(self, in_channel, depth, stride):
+        super(BasicBlockIR, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = pool_layer(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+                BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
+            BatchNorm2d(depth),
+            PReLU(depth),
+            Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
+            BatchNorm2d(depth))
+
+    def forward(self, x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+
+        return res + shortcut
+
+
+class bottleneck_IR_SE(Module):
+    def __init__(self, in_channel, depth, stride):
+        super(bottleneck_IR_SE, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = pool_layer(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+                BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
+            PReLU(depth),
+            Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
+            BatchNorm2d(depth),
+            SEModule(depth, 16)
+        )
+
+    def forward(self, x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+
+        return res + shortcut
+
+
+class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
+    '''A named tuple describing a ResNet block.'''
+
+
+def get_block(in_channel, depth, num_units, stride=2):
+
+    return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]
+
+
+def get_blocks(num_layers):
+
+    if num_layers == 8:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+        ]
+    elif num_layers == 16:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=4),
+        ]
+    elif num_layers == 34:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=4),
+            get_block(in_channel=128, depth=256, num_units=6),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+    elif num_layers == 44:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=4),
+            get_block(in_channel=128, depth=256, num_units=14),
+        ]
+    elif num_layers == 50:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),    # (B, 64, 56, 56)
+            get_block(in_channel=64, depth=128, num_units=4),   # (B, 128, 28, 28)
+            get_block(in_channel=128, depth=256, num_units=14),   # (B, 256, 14, 14)
+            get_block(in_channel=256, depth=512, num_units=3)     # (B, 512, 7, 7)
+        ]
+    elif num_layers == 100:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=13),
+            get_block(in_channel=128, depth=256, num_units=30),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+    elif num_layers == 152:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=8),
+            get_block(in_channel=128, depth=256, num_units=36),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+
+    return blocks
+
+
+@BACKBONES.register_module()
+class IRSE(BaseBackbone):
+    def __init__(self, input_size, num_layers, mode='ir', with_head=False, pretrained=None, return_index=(0, 1, 2),
+                return_type='Tuple'):
+        super().__init__()
+        assert input_size[0] in [112, 224], "input_size should be [112, 112] or [224, 224]"
+        assert num_layers in [0, 8, 16, 34, 44, 50, 100, 152], "num_layers should be 50, 100 or 152"
+        assert mode in ['ir', 'ir_se'], "mode should be ir or ir_se"
+        assert return_type in ('Tensor', 'Tuple'), 'return_tensor must be Tensor or Tuple'
+        if len(return_index) > 1:
+            assert return_type == 'Tuple'
+        self.num_layers = num_layers
+        self.return_index = return_index
+        self.return_type = return_type
+        if num_layers == 0:
+            return
+        self.with_head = with_head
+        blocks = get_blocks(num_layers)
+        if mode == 'ir':
+            if num_layers == 34:
+                warnings.warn('Using the IR_34 version from TFace')
+                unit_module = BasicBlockIR
+            else:
+                unit_module = bottleneck_IR
+        elif mode == 'ir_se':
+            unit_module = bottleneck_IR_SE
+        self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
+                                      BatchNorm2d(64),
+                                      PReLU(64))
+        if with_head:
+            if input_size[0] == 112:
+                self.output_layer = Sequential(BatchNorm2d(512),
+                                            Dropout(),
+                                            Flatten(),
+                                            Linear(512 * 7 * 7, 512),
+                                            BatchNorm1d(512))
+            else:
+                self.output_layer = Sequential(BatchNorm2d(512),
+                                            Dropout(),
+                                            Flatten(),
+                                            Linear(512 * 14 * 14, 512),
+                                            BatchNorm1d(512))
+
+        modules = []
+        max_stage = max(return_index)
+        for block in blocks[:max_stage+1]:
+            block_module = []
+            for bottleneck in block:
+                block_module.append(
+                    unit_module(bottleneck.in_channel,
+                                bottleneck.depth,
+                                bottleneck.stride))
+            modules.append(Sequential(*block_module))
+        # self.body = Sequential(*modules)
+        self.body = nn.ModuleList(modules)
+
+        self._initialize_weights()
+        if pretrained:
+            self.init_weights(pretrained)
+
+    def init_weights(self, pretrained):
+        logger = get_root_logger()
+        logger.warning(f'{self.__class__.__name__} load pretrain from {pretrained}')
+        state_dict = torch.load(pretrained, map_location='cpu')
+        if 'state_dict' in state_dict:
+            state_dict = state_dict['state_dict']
+
+        stage_unit_nums = {
+            34: (3, 4, 6, 3),
+            50: (3, 4, 14, 3),
+        }
+        stage_num = stage_unit_nums[self.num_layers]
+
+        new_state_dict = dict()
+        for k, v in state_dict.items():
+            if k.startswith('body.'):
+                index = int(k.split('.')[1])
+                if 0 <= index < stage_num[0]:
+                    new_key = k.replace('body.', 'body.0.')
+                elif stage_num[0] <= index < sum(stage_num[:2]):
+                    new_key = f"body.1.{index-sum(stage_num[:1])}.{'.'.join(k.split('.')[2:])}"
+                elif sum(stage_num[:2]) <= index < sum(stage_num[:3]):
+                    new_key = f"body.2.{index-sum(stage_num[:2])}.{'.'.join(k.split('.')[2:])}"
+                else:
+                    new_key = k
+            else:
+                new_key = k
+            new_state_dict[new_key] = v
+
+        load_state_dict(self, new_state_dict, strict=False, logger=logger)
+
+    def forward(self, x):
+        if self.num_layers == 0:
+            return x
+        x = self.input_layer(x)
+        output = []
+        return_index = set(self.return_index)
+        for index, m in enumerate(self.body):
+            x = m(x)
+            if index in return_index:
+                output.append(x)        
+        if self.with_head:
+            x = self.output_layer(x)
+        if self.return_type == 'Tensor':
+            return output[0]
+        else:
+            return tuple(output)
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_uniform_(m.weight.data)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm1d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight.data)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+    
+
+    def _freeze_stages(self):
+        if self.frozen_blocks > 0:
+            print(f'IRSE freeze the first {self.frozen_blocks} blocks, it has {len(self.body)} blocks ')
+            self.input_layer.eval()
+            print('in freeze', self.input_layer[1].training)
+            for param in self.input_layer.parameters():
+                param.requires_grad = False
+        
+        for i in range(self.frozen_blocks):
+            m = self.body[i]
+            m.eval()
+            for param in m.parameters():
+                param.requires_grad = False
+
+def IR_50(input_size):
+    """Constructs a ir-50 model.
+    """
+    model = Backbone(input_size, 50, 'ir')
+
+    return model
+
+
+def IR_101(input_size):
+    """Constructs a ir-101 model.
+    """
+    model = Backbone(input_size, 100, 'ir')
+
+    return model
+
+
+def IR_152(input_size):
+    """Constructs a ir-152 model.
+    """
+    model = Backbone(input_size, 152, 'ir')
+
+    return model
+
+
+def IR_SE_50(input_size):
+    """Constructs a ir_se-50 model.
+    """
+    model = Backbone(input_size, 50, 'ir_se')
+
+    return model
+
+
+def IR_SE_101(input_size):
+    """Constructs a ir_se-101 model.
+    """
+    model = Backbone(input_size, 100, 'ir_se')
+
+    return model
+
+
+def IR_SE_152(input_size):
+    """Constructs a ir_se-152 model.
+    """
+    model = Backbone(input_size, 152, 'ir_se')
+
+    return model

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/irse_nopadding.py

@@ -0,0 +1,366 @@
+import warnings
+import torch.nn as nn
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, ReLU, Sigmoid, Dropout, MaxPool2d, \
+    AdaptiveAvgPool2d, Sequential, Module
+from collections import namedtuple
+import torch
+from mmcv.runner import load_state_dict
+from mmcls.utils import get_root_logger
+
+from ..builder import BACKBONES
+from .base_backbone import BaseBackbone
+
+# Support: ['IR_50', 'IR_101', 'IR_152', 'IR_SE_50', 'IR_SE_101', 'IR_SE_152']
+
+class Flatten(Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+
+def l2_norm(input, axis=1):
+    norm = torch.norm(input, 2, axis, True)
+    output = torch.div(input, norm)
+
+    return output
+
+
+class SEModule(Module):
+    def __init__(self, channels, reduction):
+        super(SEModule, self).__init__()
+        self.avg_pool = AdaptiveAvgPool2d(1)
+        self.fc1 = Conv2d(
+            channels, channels // reduction, kernel_size=1, padding=0, bias=False)
+
+        nn.init.xavier_uniform_(self.fc1.weight.data)
+
+        self.relu = ReLU(inplace=True)
+        self.fc2 = Conv2d(
+            channels // reduction, channels, kernel_size=1, padding=0, bias=False)
+
+        self.sigmoid = Sigmoid()
+
+    def forward(self, x):
+        module_input = x
+        x = self.avg_pool(x)
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        x = self.sigmoid(x)
+
+        return module_input * x
+
+
+class bottleneck_IR(Module):
+    def __init__(self, in_channel, depth, stride):
+        super(bottleneck_IR, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = MaxPool2d(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+                BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
+            PReLU(depth),
+            Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
+            BatchNorm2d(depth))
+
+    def forward(self, x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+
+        return res + shortcut
+
+
+class BasicBlockIR(Module):
+    """
+    BasicBlock for IRNet, stolen from TFace
+    https://github.com/Tencent/TFace/blob/d57fd8d9ce9502240921f0998c57f84afa7eaeaa/torchkit/backbone/model_irse.py#L16
+    Add a BatchNorm2d after the first Conv2d
+    """
+    def __init__(self, in_channel, depth, stride):
+        super(BasicBlockIR, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = MaxPool2d(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+                BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3, 3), (1, 1), 0, bias=False),
+            BatchNorm2d(depth),
+            PReLU(depth),
+            Conv2d(depth, depth, (3, 3), stride, 0, bias=False),
+            BatchNorm2d(depth))
+
+    def forward(self, x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+
+        return res + shortcut
+
+
+class bottleneck_IR_SE(Module):
+    def __init__(self, in_channel, depth, stride):
+        super(bottleneck_IR_SE, self).__init__()
+        if in_channel == depth:
+            self.shortcut_layer = MaxPool2d(1, stride)
+        else:
+            self.shortcut_layer = Sequential(
+                Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+                BatchNorm2d(depth))
+        self.res_layer = Sequential(
+            BatchNorm2d(in_channel),
+            Conv2d(in_channel, depth, (3, 3), (1, 1), 0, bias=False),
+            PReLU(depth),
+            Conv2d(depth, depth, (3, 3), stride, 0, bias=False),
+            BatchNorm2d(depth),
+            SEModule(depth, 16)
+        )
+
+    def forward(self, x):
+        shortcut = self.shortcut_layer(x)
+        res = self.res_layer(x)
+
+        return res + shortcut
+
+
+class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
+    '''A named tuple describing a ResNet block.'''
+
+
+def get_block(in_channel, depth, num_units, stride=2):
+
+    return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]
+
+
+def get_blocks(num_layers):
+
+    if num_layers == 8:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+        ]
+    elif num_layers == 16:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=4),
+        ]
+    elif num_layers == 34:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=4),
+            get_block(in_channel=128, depth=256, num_units=6),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+    elif num_layers == 44:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=4),
+            get_block(in_channel=128, depth=256, num_units=14),
+        ]
+    elif num_layers == 50:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),    # (B, 64, 56, 56)
+            get_block(in_channel=64, depth=128, num_units=4),   # (B, 128, 28, 28)
+            get_block(in_channel=128, depth=256, num_units=14),   # (B, 256, 14, 14)
+            get_block(in_channel=256, depth=512, num_units=3)     # (B, 512, 7, 7)
+        ]
+    elif num_layers == 100:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=13),
+            get_block(in_channel=128, depth=256, num_units=30),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+    elif num_layers == 152:
+        blocks = [
+            get_block(in_channel=64, depth=64, num_units=3),
+            get_block(in_channel=64, depth=128, num_units=8),
+            get_block(in_channel=128, depth=256, num_units=36),
+            get_block(in_channel=256, depth=512, num_units=3)
+        ]
+
+    return blocks
+
+
+@BACKBONES.register_module()
+class IRSENoPadding(BaseBackbone):
+    def __init__(self, input_size, num_layers, mode='ir', with_head=False, pretrained=None, return_index=(0, 1, 2)):
+        super().__init__()
+        assert input_size[0] in [112, 224], "input_size should be [112, 112] or [224, 224]"
+        assert num_layers in [0, 8, 16, 34, 44, 50, 100, 152], "num_layers should be 50, 100 or 152"
+        assert mode in ['ir', 'ir_se'], "mode should be ir or ir_se"
+        self.num_layers = num_layers
+        self.return_index = return_index
+        if num_layers == 0:
+            return
+        self.with_head = with_head
+        blocks = get_blocks(num_layers)
+        if mode == 'ir':
+            if num_layers == 34:
+                warnings.warn('Using the IR_34 version from TFace')
+                unit_module = BasicBlockIR
+            else:
+                unit_module = bottleneck_IR
+        elif mode == 'ir_se':
+            unit_module = bottleneck_IR_SE
+        self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 0, bias=False),
+                                      BatchNorm2d(64),
+                                      PReLU(64))
+        if with_head:
+            if input_size[0] == 112:
+                self.output_layer = Sequential(BatchNorm2d(512),
+                                            Dropout(),
+                                            Flatten(),
+                                            Linear(512 * 7 * 7, 512),
+                                            BatchNorm1d(512))
+            else:
+                self.output_layer = Sequential(BatchNorm2d(512),
+                                            Dropout(),
+                                            Flatten(),
+                                            Linear(512 * 14 * 14, 512),
+                                            BatchNorm1d(512))
+
+        modules = []
+        max_stage = max(return_index)
+        for block in blocks[:max_stage+1]:
+            block_module = []
+            for bottleneck in block:
+                block_module.append(
+                    unit_module(bottleneck.in_channel,
+                                bottleneck.depth,
+                                bottleneck.stride))
+            modules.append(Sequential(*block_module))
+        # self.body = Sequential(*modules)
+        self.body = nn.ModuleList(modules)
+
+        self._initialize_weights()
+        if pretrained:
+            self.init_weights(pretrained)
+
+    def init_weights(self, pretrained):
+        logger = get_root_logger()
+        logger.warning(f'{self.__class__.__name__} load pretrain from {pretrained}')
+        state_dict = torch.load(pretrained, map_location='cpu')
+        if 'state_dict' in state_dict:
+            state_dict = state_dict['state_dict']
+
+        stage_unit_nums = {
+            34: (3, 4, 6, 3),
+            50: (3, 4, 14, 3),
+        }
+        stage_num = stage_unit_nums[self.num_layers]
+
+        new_state_dict = dict()
+        for k, v in state_dict.items():
+            if k.startswith('body.'):
+                index = int(k.split('.')[1])
+                if 0 <= index < stage_num[0]:
+                    new_key = k.replace('body.', 'body.0.')
+                elif stage_num[0] <= index < sum(stage_num[:2]):
+                    new_key = f"body.1.{index-sum(stage_num[:1])}.{'.'.join(k.split('.')[2:])}"
+                elif sum(stage_num[:2]) <= index < sum(stage_num[:3]):
+                    new_key = f"body.2.{index-sum(stage_num[:2])}.{'.'.join(k.split('.')[2:])}"
+                else:
+                    new_key = k
+            else:
+                new_key = k
+            new_state_dict[new_key] = v
+
+        load_state_dict(self, new_state_dict, strict=False, logger=logger)
+
+    def forward(self, x):
+        if self.num_layers == 0:
+            return x
+        x = self.input_layer(x)
+        output = []
+        return_index = set(self.return_index)
+        for index, m in enumerate(self.body):
+            x = m(x)
+            if index in return_index:
+                output.append(x)        
+        if self.with_head:
+            x = self.output_layer(x)
+
+        return output
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_uniform_(m.weight.data)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm1d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight.data)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+    
+
+    def _freeze_stages(self):
+        if self.frozen_blocks > 0:
+            print(f'IRSE freeze the first {self.frozen_blocks} blocks, it has {len(self.body)} blocks ')
+            self.input_layer.eval()
+            print('in freeze', self.input_layer[1].training)
+            for param in self.input_layer.parameters():
+                param.requires_grad = False
+        
+        for i in range(self.frozen_blocks):
+            m = self.body[i]
+            m.eval()
+            for param in m.parameters():
+                param.requires_grad = False
+
+def IR_50(input_size):
+    """Constructs a ir-50 model.
+    """
+    model = Backbone(input_size, 50, 'ir')
+
+    return model
+
+
+def IR_101(input_size):
+    """Constructs a ir-101 model.
+    """
+    model = Backbone(input_size, 100, 'ir')
+
+    return model
+
+
+def IR_152(input_size):
+    """Constructs a ir-152 model.
+    """
+    model = Backbone(input_size, 152, 'ir')
+
+    return model
+
+
+def IR_SE_50(input_size):
+    """Constructs a ir_se-50 model.
+    """
+    model = Backbone(input_size, 50, 'ir_se')
+
+    return model
+
+
+def IR_SE_101(input_size):
+    """Constructs a ir_se-101 model.
+    """
+    model = Backbone(input_size, 100, 'ir_se')
+
+    return model
+
+
+def IR_SE_152(input_size):
+    """Constructs a ir_se-152 model.
+    """
+    model = Backbone(input_size, 152, 'ir_se')
+
+    return model

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/lenet.py

@@ -0,0 +1,41 @@
+import torch.nn as nn
+
+from ..builder import BACKBONES
+from .base_backbone import BaseBackbone
+
+
+@BACKBONES.register_module()
+class LeNet5(BaseBackbone):
+    """`LeNet5 <https://en.wikipedia.org/wiki/LeNet>`_ backbone.
+
+    The input for LeNet-5 is a 32×32 grayscale image.
+
+    Args:
+        num_classes (int): number of classes for classification.
+            The default value is -1, which uses the backbone as
+            a feature extractor without the top classifier.
+    """
+
+    def __init__(self, num_classes=-1):
+        super(LeNet5, self).__init__()
+        self.num_classes = num_classes
+        self.features = nn.Sequential(
+            nn.Conv2d(1, 6, kernel_size=5, stride=1), nn.Tanh(),
+            nn.AvgPool2d(kernel_size=2),
+            nn.Conv2d(6, 16, kernel_size=5, stride=1), nn.Tanh(),
+            nn.AvgPool2d(kernel_size=2),
+            nn.Conv2d(16, 120, kernel_size=5, stride=1), nn.Tanh())
+        if self.num_classes > 0:
+            self.classifier = nn.Sequential(
+                nn.Linear(120, 84),
+                nn.Tanh(),
+                nn.Linear(84, num_classes),
+            )
+
+    def forward(self, x):
+
+        x = self.features(x)
+        if self.num_classes > 0:
+            x = self.classifier(x.squeeze())
+
+        return x

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/mobilefacenet.py

@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, Module, Sequential
+
+from ..builder import BACKBONES
+from .base_backbone import BaseBackbone
+
+
+class Flatten(Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+
+def l2_norm(input,axis=1):
+    norm = torch.norm(input,2,axis,True)
+    output = torch.div(input, norm)
+    return output
+
+
+class Conv_block(Module):
+    def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
+        super(Conv_block, self).__init__()
+        self.conv = Conv2d(in_c, out_channels=out_c, kernel_size=kernel, groups=groups, stride=stride, padding=padding, bias=False)
+        self.bn = BatchNorm2d(out_c)
+        self.prelu = PReLU(out_c)
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.prelu(x)
+        return x
+
+class Linear_block(Module):
+    def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
+        super(Linear_block, self).__init__()
+        self.conv = Conv2d(in_c, out_channels=out_c, kernel_size=kernel, groups=groups, stride=stride, padding=padding, bias=False)
+        self.bn = BatchNorm2d(out_c)
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+class Depth_Wise(Module):
+     def __init__(self, in_c, out_c, residual = False, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=1):
+        super(Depth_Wise, self).__init__()
+        self.conv = Conv_block(in_c, out_c=groups, kernel=(1, 1), padding=(0, 0), stride=(1, 1))
+        self.conv_dw = Conv_block(groups, groups, groups=groups, kernel=kernel, padding=padding, stride=stride)
+        self.project = Linear_block(groups, out_c, kernel=(1, 1), padding=(0, 0), stride=(1, 1))
+        self.residual = residual
+     def forward(self, x):
+        short_cut = None
+        if self.residual:
+            short_cut = x
+        x = self.conv(x)
+        x = self.conv_dw(x)
+        x = self.project(x)
+        if self.residual:
+            output = short_cut + x
+        else:
+            output = x
+        return output
+
+class Residual(Module):
+    def __init__(self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1)):
+        super(Residual, self).__init__()
+        modules = []
+        for _ in range(num_block):
+            modules.append(Depth_Wise(c, c, residual=True, kernel=kernel, padding=padding, stride=stride, groups=groups))
+        self.model = Sequential(*modules)
+    def forward(self, x):
+        return self.model(x)
+
+
+@BACKBONES.register_module()
+class MobileFaceNet(BaseBackbone):
+    """`MobileFaceNet <https://github.com/TreB1eN/InsightFace_Pytorch>`_ backbone.
+
+    The input for MobileFaceNet is a 224x224 RGB image.
+
+    Args:
+        num_classes (int): number of classes for classification.
+            The default value is -1, which uses the backbone as
+            a feature extractor without the top classifier.
+    """
+    def __init__(self, return_stages=[2], pretrained=None):
+        super().__init__()
+        self.conv1 = Conv_block(3, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1))
+
+        self.conv2_dw = Conv_block(64, 64, kernel=(3, 3), stride=(1, 1), padding=(1, 1), groups=64)
+
+        self.conv_23 = Depth_Wise(64, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=128)
+        self.conv_3 = Residual(64, num_block=4, groups=128, kernel=(3, 3), stride=(1, 1), padding=(1, 1))
+
+        self.conv_34 = Depth_Wise(64, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=256)
+        self.conv_4 = Residual(128, num_block=6, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1))
+
+        # self.conv_45 = Depth_Wise(128, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=512)
+        # self.conv_5 = Residual(128, num_block=2, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1))
+
+        # self.conv_6_sep = Conv_block(128, 512, kernel=(1, 1), stride=(1, 1), padding=(0, 0))
+        # self.conv_6_dw = Linear_block(512, 512, groups=512, kernel=(7,7), stride=(1, 1), padding=(0, 0))
+        # self.conv_6_flatten = Flatten()
+        # self.linear = Linear(512, embedding_size, bias=False)
+        # self.bn = BatchNorm1d(embedding_size)
+        self.return_stages = set(return_stages)
+        if pretrained:
+            self.init_weights(pretrained)
+    
+    def forward(self, x):
+        output = []
+        out = self.conv1(x)
+        out = self.conv2_dw(out)
+        if 0 in self.return_stages:
+            output.append(out)
+        
+        out = self.conv_23(out)
+        out = self.conv_3(out)
+        if 1 in self.return_stages:
+            output.append(out)
+        
+        out = self.conv_34(out)
+        out = self.conv_4(out)
+        if 2 in self.return_stages:
+            output.append(out)
+        # out = self.conv_45(out)
+        # out = self.conv_5(out)
+        # out = self.conv_6_sep(out)
+        # out = self.conv_6_dw(out)
+        # out = self.conv_6_flatten(out)
+        # out = self.linear(out)
+        # out = self.bn(out)
+        # print(out.shape)
+        # return out
+        return output

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/mobilenet_v2.py

@@ -0,0 +1,273 @@
+import logging
+
+import torch.nn as nn
+import torch.utils.checkpoint as cp
+from mmcv.cnn import ConvModule, constant_init, kaiming_init
+from mmcv.runner import load_checkpoint
+from torch.nn.modules.batchnorm import _BatchNorm
+
+from mmcls.models.utils import make_divisible
+from ..builder import BACKBONES
+from .base_backbone import BaseBackbone
+
+
+class InvertedResidual(nn.Module):
+    """InvertedResidual block for MobileNetV2.
+
+    Args:
+        in_channels (int): The input channels of the InvertedResidual block.
+        out_channels (int): The output channels of the InvertedResidual block.
+        stride (int): Stride of the middle (first) 3x3 convolution.
+        expand_ratio (int): adjusts number of channels of the hidden layer
+            in InvertedResidual by this amount.
+        conv_cfg (dict): Config dict for convolution layer.
+            Default: None, which means using conv2d.
+        norm_cfg (dict): Config dict for normalization layer.
+            Default: dict(type='BN').
+        act_cfg (dict): Config dict for activation layer.
+            Default: dict(type='ReLU6').
+        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
+            memory while slowing down the training speed. Default: False.
+
+    Returns:
+        Tensor: The output tensor
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 stride,
+                 expand_ratio,
+                 conv_cfg=None,
+                 norm_cfg=dict(type='BN'),
+                 act_cfg=dict(type='ReLU6'),
+                 with_cp=False):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2], f'stride must in [1, 2]. ' \
+            f'But received {stride}.'
+        self.with_cp = with_cp
+        self.use_res_connect = self.stride == 1 and in_channels == out_channels
+        hidden_dim = int(round(in_channels * expand_ratio))
+
+        layers = []
+        if expand_ratio != 1:
+            layers.append(
+                ConvModule(
+                    in_channels=in_channels,
+                    out_channels=hidden_dim,
+                    kernel_size=1,
+                    conv_cfg=conv_cfg,
+                    norm_cfg=norm_cfg,
+                    act_cfg=act_cfg))
+        layers.extend([
+            ConvModule(
+                in_channels=hidden_dim,
+                out_channels=hidden_dim,
+                kernel_size=3,
+                stride=stride,
+                padding=1,
+                groups=hidden_dim,
+                conv_cfg=conv_cfg,
+                norm_cfg=norm_cfg,
+                act_cfg=act_cfg),
+            ConvModule(
+                in_channels=hidden_dim,
+                out_channels=out_channels,
+                kernel_size=1,
+                conv_cfg=conv_cfg,
+                norm_cfg=norm_cfg,
+                act_cfg=None)
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+
+        def _inner_forward(x):
+            if self.use_res_connect:
+                return x + self.conv(x)
+            else:
+                return self.conv(x)
+
+        if self.with_cp and x.requires_grad:
+            out = cp.checkpoint(_inner_forward, x)
+        else:
+            out = _inner_forward(x)
+
+        return out
+
+
+@BACKBONES.register_module()
+class MobileNetV2(BaseBackbone):
+    """MobileNetV2 backbone.
+
+    Args:
+        widen_factor (float): Width multiplier, multiply number of
+            channels in each layer by this amount. Default: 1.0.
+        out_indices (None or Sequence[int]): Output from which stages.
+            Default: (7, ).
+        frozen_stages (int): Stages to be frozen (all param fixed).
+            Default: -1, which means not freezing any parameters.
+        conv_cfg (dict): Config dict for convolution layer.
+            Default: None, which means using conv2d.
+        norm_cfg (dict): Config dict for normalization layer.
+            Default: dict(type='BN').
+        act_cfg (dict): Config dict for activation layer.
+            Default: dict(type='ReLU6').
+        norm_eval (bool): Whether to set norm layers to eval mode, namely,
+            freeze running stats (mean and var). Note: Effect on Batch Norm
+            and its variants only. Default: False.
+        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
+            memory while slowing down the training speed. Default: False.
+    """
+
+    # Parameters to build layers. 4 parameters are needed to construct a
+    # layer, from left to right: expand_ratio, channel, num_blocks, stride.
+    arch_settings = [[1, 16, 1, 1], [6, 24, 2, 2], [6, 32, 3, 2],
+                     [6, 64, 4, 2], [6, 96, 3, 1], [6, 160, 3, 2],
+                     [6, 320, 1, 1]]
+
+    def __init__(self,
+                 widen_factor=1.,
+                 out_indices=(7, ),
+                 frozen_stages=-1,
+                 conv_cfg=None,
+                 norm_cfg=dict(type='BN'),
+                 act_cfg=dict(type='ReLU6'),
+                 norm_eval=False,
+                 with_cp=False):
+        super(MobileNetV2, self).__init__()
+        self.widen_factor = widen_factor
+        self.out_indices = out_indices
+        for index in out_indices:
+            if index not in range(0, 8):
+                raise ValueError('the item in out_indices must in '
+                                 f'range(0, 8). But received {index}')
+
+        if frozen_stages not in range(-1, 8):
+            raise ValueError('frozen_stages must be in range(-1, 8). '
+                             f'But received {frozen_stages}')
+        self.out_indices = out_indices
+        self.frozen_stages = frozen_stages
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.act_cfg = act_cfg
+        self.norm_eval = norm_eval
+        self.with_cp = with_cp
+
+        self.in_channels = make_divisible(32 * widen_factor, 8)
+
+        self.conv1 = ConvModule(
+            in_channels=3,
+            out_channels=self.in_channels,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            conv_cfg=self.conv_cfg,
+            norm_cfg=self.norm_cfg,
+            act_cfg=self.act_cfg)
+
+        self.layers = []
+
+        for i, layer_cfg in enumerate(self.arch_settings):
+            expand_ratio, channel, num_blocks, stride = layer_cfg
+            out_channels = make_divisible(channel * widen_factor, 8)
+            inverted_res_layer = self.make_layer(
+                out_channels=out_channels,
+                num_blocks=num_blocks,
+                stride=stride,
+                expand_ratio=expand_ratio)
+            layer_name = f'layer{i + 1}'
+            self.add_module(layer_name, inverted_res_layer)
+            self.layers.append(layer_name)
+
+        if widen_factor > 1.0:
+            self.out_channel = int(1280 * widen_factor)
+        else:
+            self.out_channel = 1280
+
+        layer = ConvModule(
+            in_channels=self.in_channels,
+            out_channels=self.out_channel,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            conv_cfg=self.conv_cfg,
+            norm_cfg=self.norm_cfg,
+            act_cfg=self.act_cfg)
+        self.add_module('conv2', layer)
+        self.layers.append('conv2')
+
+    def make_layer(self, out_channels, num_blocks, stride, expand_ratio):
+        """ Stack InvertedResidual blocks to build a layer for MobileNetV2.
+
+        Args:
+            out_channels (int): out_channels of block.
+            num_blocks (int): number of blocks.
+            stride (int): stride of the first block. Default: 1
+            expand_ratio (int): Expand the number of channels of the
+                hidden layer in InvertedResidual by this ratio. Default: 6.
+        """
+        layers = []
+        for i in range(num_blocks):
+            if i >= 1:
+                stride = 1
+            layers.append(
+                InvertedResidual(
+                    self.in_channels,
+                    out_channels,
+                    stride,
+                    expand_ratio=expand_ratio,
+                    conv_cfg=self.conv_cfg,
+                    norm_cfg=self.norm_cfg,
+                    act_cfg=self.act_cfg,
+                    with_cp=self.with_cp))
+            self.in_channels = out_channels
+
+        return nn.Sequential(*layers)
+
+    def init_weights(self, pretrained=None):
+        if isinstance(pretrained, str):
+            logger = logging.getLogger()
+            load_checkpoint(self, pretrained, strict=False, logger=logger)
+        elif pretrained is None:
+            for m in self.modules():
+                if isinstance(m, nn.Conv2d):
+                    kaiming_init(m)
+                elif isinstance(m, (_BatchNorm, nn.GroupNorm)):
+                    constant_init(m, 1)
+        else:
+            raise TypeError('pretrained must be a str or None')
+
+    def forward(self, x):
+        x = self.conv1(x)
+
+        outs = []
+        for i, layer_name in enumerate(self.layers):
+            layer = getattr(self, layer_name)
+            x = layer(x)
+            if i in self.out_indices:
+                outs.append(x)
+
+        if len(outs) == 1:
+            return outs[0]
+        else:
+            return tuple(outs)
+
+    def _freeze_stages(self):
+        if self.frozen_stages >= 0:
+            for param in self.conv1.parameters():
+                param.requires_grad = False
+        for i in range(1, self.frozen_stages + 1):
+            layer = getattr(self, f'layer{i}')
+            layer.eval()
+            for param in layer.parameters():
+                param.requires_grad = False
+
+    def train(self, mode=True):
+        super(MobileNetV2, self).train(mode)
+        self._freeze_stages()
+        if mode and self.norm_eval:
+            for m in self.modules():
+                if isinstance(m, _BatchNorm):
+                    m.eval()

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/mobilenet_v3.py

@@ -0,0 +1,186 @@
+import logging
+
+import torch.nn as nn
+from mmcv.cnn import ConvModule, constant_init, kaiming_init
+from mmcv.runner import load_checkpoint
+from torch.nn.modules.batchnorm import _BatchNorm
+
+from ..builder import BACKBONES
+from ..utils import InvertedResidual
+from .base_backbone import BaseBackbone
+
+
+@BACKBONES.register_module()
+class MobileNetv3(BaseBackbone):
+    """ MobileNetv3 backbone
+
+    Args:
+        arch (str): Architechture of mobilnetv3, from {small, big}.
+            Default: small.
+        conv_cfg (dict): Config dict for convolution layer.
+            Default: None, which means using conv2d.
+        norm_cfg (dict): Config dict for normalization layer.
+            Default: dict(type='BN').
+        out_indices (None or Sequence[int]): Output from which stages.
+            Default: (10, ), which means output tensors from final stage.
+        frozen_stages (int): Stages to be frozen (all param fixed).
+            Defualt: -1, which means not freezing any parameters.
+        norm_eval (bool): Whether to set norm layers to eval mode, namely,
+            freeze running stats (mean and var). Note: Effect on Batch Norm
+            and its variants only. Default: False.
+        with_cp (bool): Use checkpoint or not. Using checkpoint will save
+            some memory while slowing down the training speed.
+            Defualt: False.
+    """
+    # Parameters to build each block:
+    #     [kernel size, mid channels, out channels, with_se, act type, stride]
+    arch_settings = {
+        'small': [[3, 16, 16, True, 'ReLU', 2],
+                  [3, 72, 24, False, 'ReLU', 2],
+                  [3, 88, 24, False, 'ReLU', 1],
+                  [5, 96, 40, True, 'HSwish', 2],
+                  [5, 240, 40, True, 'HSwish', 1],
+                  [5, 240, 40, True, 'HSwish', 1],
+                  [5, 120, 48, True, 'HSwish', 1],
+                  [5, 144, 48, True, 'HSwish', 1],
+                  [5, 288, 96, True, 'HSwish', 2],
+                  [5, 576, 96, True, 'HSwish', 1],
+                  [5, 576, 96, True, 'HSwish', 1]],
+        'big': [[3, 16, 16, False, 'ReLU', 1],
+                [3, 64, 24, False, 'ReLU', 2],
+                [3, 72, 24, False, 'ReLU', 1],
+                [5, 72, 40, True, 'ReLU', 2],
+                [5, 120, 40, True, 'ReLU', 1],
+                [5, 120, 40, True, 'ReLU', 1],
+                [3, 240, 80, False, 'HSwish', 2],
+                [3, 200, 80, False, 'HSwish', 1],
+                [3, 184, 80, False, 'HSwish', 1],
+                [3, 184, 80, False, 'HSwish', 1],
+                [3, 480, 112, True, 'HSwish', 1],
+                [3, 672, 112, True, 'HSwish', 1],
+                [5, 672, 160, True, 'HSwish', 1],
+                [5, 672, 160, True, 'HSwish', 2],
+                [5, 960, 160, True, 'HSwish', 1]]
+    }  # yapf: disable
+
+    def __init__(self,
+                 arch='small',
+                 conv_cfg=None,
+                 norm_cfg=dict(type='BN'),
+                 out_indices=(10, ),
+                 frozen_stages=-1,
+                 norm_eval=False,
+                 with_cp=False):
+        super(MobileNetv3, self).__init__()
+        assert arch in self.arch_settings
+        for index in out_indices:
+            if index not in range(0, len(self.arch_settings[arch])):
+                raise ValueError('the item in out_indices must in '
+                                 f'range(0, {len(self.arch_settings[arch])}). '
+                                 f'But received {index}')
+
+        if frozen_stages not in range(-1, len(self.arch_settings[arch])):
+            raise ValueError('frozen_stages must be in range(-1, '
+                             f'{len(self.arch_settings[arch])}). '
+                             f'But received {frozen_stages}')
+        self.out_indices = out_indices
+        self.frozen_stages = frozen_stages
+        self.arch = arch
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.out_indices = out_indices
+        self.frozen_stages = frozen_stages
+        self.norm_eval = norm_eval
+        self.with_cp = with_cp
+
+        self.in_channels = 16
+        self.conv1 = ConvModule(
+            in_channels=3,
+            out_channels=self.in_channels,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            conv_cfg=conv_cfg,
+            norm_cfg=norm_cfg,
+            act_cfg=dict(type='HSwish'))
+
+        self.layers = self._make_layer()
+        self.feat_dim = self.arch_settings[arch][-1][2]
+
+    def _make_layer(self):
+        layers = []
+        layer_setting = self.arch_settings[self.arch]
+        for i, params in enumerate(layer_setting):
+            (kernel_size, mid_channels, out_channels, with_se, act,
+             stride) = params
+            if with_se:
+                se_cfg = dict(
+                    channels=mid_channels,
+                    ratio=4,
+                    act_cfg=(dict(type='ReLU'), dict(type='HSigmoid')))
+            else:
+                se_cfg = None
+
+            layer = InvertedResidual(
+                in_channels=self.in_channels,
+                out_channels=out_channels,
+                mid_channels=mid_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                se_cfg=se_cfg,
+                with_expand_conv=True,
+                conv_cfg=self.conv_cfg,
+                norm_cfg=self.norm_cfg,
+                act_cfg=dict(type=act),
+                with_cp=self.with_cp)
+            self.in_channels = out_channels
+            layer_name = 'layer{}'.format(i + 1)
+            self.add_module(layer_name, layer)
+            layers.append(layer_name)
+        return layers
+
+    def init_weights(self, pretrained=None):
+        if isinstance(pretrained, str):
+            logger = logging.getLogger()
+            load_checkpoint(self, pretrained, strict=False, logger=logger)
+        elif pretrained is None:
+            for m in self.modules():
+                if isinstance(m, nn.Conv2d):
+                    kaiming_init(m)
+                elif isinstance(m, nn.BatchNorm2d):
+                    constant_init(m, 1)
+        else:
+            raise TypeError('pretrained must be a str or None')
+
+    def forward(self, x):
+        x = self.conv1(x)
+
+        outs = []
+        for i, layer_name in enumerate(self.layers):
+            layer = getattr(self, layer_name)
+            x = layer(x)
+            if i in self.out_indices:
+                outs.append(x)
+
+        if len(outs) == 1:
+            return outs[0]
+        else:
+            return tuple(outs)
+
+    def _freeze_stages(self):
+        if self.frozen_stages >= 0:
+            for param in self.conv1.parameters():
+                param.requires_grad = False
+        for i in range(1, self.frozen_stages + 1):
+            layer = getattr(self, f'layer{i}')
+            layer.eval()
+            for param in layer.parameters():
+                param.requires_grad = False
+
+    def train(self, mode=True):
+        super(MobileNetv3, self).train(mode)
+        self._freeze_stages()
+        if mode and self.norm_eval:
+            for m in self.modules():
+                if isinstance(m, _BatchNorm):
+                    m.eval()

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/modules/t2t.py

@@ -0,0 +1,57 @@
+import math
+import torch
+from torch import nn
+
+
+class Token_performer(nn.Module):
+    def __init__(self, dim, in_dim, head_cnt=1, kernel_ratio=0.5, dp1=0.1, dp2 = 0.1):
+        super().__init__()
+        self.emb = in_dim * head_cnt # we use 1, so it is no need here
+        self.kqv = nn.Linear(dim, 3 * self.emb)
+        self.dp = nn.Dropout(dp1)
+        self.proj = nn.Linear(self.emb, self.emb)
+        self.head_cnt = head_cnt
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(self.emb)
+        self.epsilon = 1e-8  # for stable in division
+
+        self.mlp = nn.Sequential(
+            nn.Linear(self.emb, 1 * self.emb),
+            nn.GELU(),
+            nn.Linear(1 * self.emb, self.emb),
+            nn.Dropout(dp2),
+        )
+
+        self.m = int(self.emb * kernel_ratio)
+        self.w = torch.randn(self.m, self.emb)
+        self.w = nn.Parameter(nn.init.orthogonal_(self.w) * math.sqrt(self.m), requires_grad=False)
+
+    def prm_exp(self, x):
+        # part of the function is borrow from https://github.com/lucidrains/performer-pytorch 
+        # and Simo Ryu (https://github.com/cloneofsimo)
+        # ==== positive random features for gaussian kernels ====
+        # x = (B, T, hs)
+        # w = (m, hs)
+        # return : x : B, T, m
+        # SM(x, y) = E_w[exp(w^T x - |x|/2) exp(w^T y - |y|/2)]
+        # therefore return exp(w^Tx - |x|/2)/sqrt(m)
+        xd = ((x * x).sum(dim=-1, keepdim=True)).repeat(1, 1, self.m) / 2
+        wtx = torch.einsum('bti,mi->btm', x.float(), self.w)
+
+        return torch.exp(wtx - xd) / math.sqrt(self.m)
+
+    def single_attn(self, x):
+        k, q, v = torch.split(self.kqv(x), self.emb, dim=-1)
+        kp, qp = self.prm_exp(k), self.prm_exp(q)  # (B, T, m), (B, T, m)
+        D = torch.einsum('bti,bi->bt', qp, kp.sum(dim=1)).unsqueeze(dim=2)  # (B, T, m) * (B, m) -> (B, T, 1)
+        kptv = torch.einsum('bin,bim->bnm', v.float(), kp)  # (B, emb, m)
+        y = torch.einsum('bti,bni->btn', qp, kptv) / (D.repeat(1, 1, self.emb) + self.epsilon)  # (B, T, emb)/Diag
+        # skip connection
+        y = v + self.dp(self.proj(y))  # same as token_transformer in T2T layer, use v as skip connection
+
+        return y
+
+    def forward(self, x):
+        x = self.single_attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/modules/vit.py

@@ -0,0 +1,197 @@
+import torch
+from torch import nn
+from mmcls.models.vit.layers import DropPath
+
+
+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
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        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)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale   # [B, head_num, token_num, token_num]
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class HeadFusionAttention(nn.Module):
+    """
+    fuse front head output add current head input as current input
+    Origin: head2_output = f(head2_input)
+    Fuse: head2_output = f(head1_output + head2_input)
+    """
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.group_number = 4 # 每个group内部并行计算
+        self.qkv = nn.ModuleList([nn.Linear(dim//self.group_number, (dim//self.group_number)*3, bias=qkv_bias) for _ in range(self.group_number)])
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        # H = self.num_heads
+        # dim = C // self.num_heads
+        g = self.group_number
+        g_dim = C // self.group_number
+        x = x.reshape((B, N, g, g_dim))
+
+        outputs = []
+        head_x = torch.zeros((B, N, g_dim), device=x.device)
+        for i in range(g):
+            # self-attention
+            current_x = x[:, :, i]
+            current_x = current_x + head_x
+            qkv = self.qkv[i](current_x).reshape(B, N, 3, g_dim)
+            qkv = qkv.permute(2, 0, 1, 3)   # [3, B, N, dim]
+            q, k, v = qkv[0], qkv[1], qkv[2] 
+            attn = (q @ k.transpose(-2, -1)) * self.scale   # [B, N, N]
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            head_x = (attn @ v)  # [B, N, d]
+            outputs.append(head_x)
+        x = torch.cat(outputs, dim=-1)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class HeadFusionAttentionV2(nn.Module):
+    """
+    V2:  one branch is origin, onther is modified
+    fuse front head output add current head input as current input
+    Origin: head2_output = f(head2_input)
+    Fuse: head2_output = f(head1_output + head2_input)
+    """
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.group_number = 2 # 每个group内部并行计算.
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.qkv2 = nn.ModuleList([nn.Linear(dim//self.group_number, (dim//self.group_number)*3, bias=qkv_bias) for _ in range(self.group_number)])
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        # original        
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale   # [B, head_num, token_num, token_num]
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        origin = (attn @ v)     # [B, H, N, dim]
+
+        # modify
+        g = self.group_number
+        g_dim = C // self.group_number
+        x = x.reshape((B, N, g, g_dim))
+        n = self.num_heads // g  # head number per group
+
+        origin = origin.transpose(1, 2).reshape((B, N, g, g_dim))
+
+        outputs = []
+        head_x = torch.zeros((B, N, g_dim), device=x.device)
+        for i in range(g):
+            # self-attention
+            current_x = x[:, :, i]
+            current_x = current_x + head_x
+            qkv = self.qkv2[i](current_x).reshape(B, N, 3, g_dim)
+            qkv = qkv.permute(2, 0, 1, 3)   # [3, B, N, dim]
+            q, k, v = qkv[0], qkv[1], qkv[2] 
+            attn = (q @ k.transpose(-2, -1)) * self.scale   # [B, N, N]
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            head_x = (attn @ v)  # [B, N, d]
+            outputs.append(head_x + origin[:, :, i])
+        x = torch.cat(outputs, dim=-1)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, head_fusion=False,
+                 ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        if head_fusion:
+            self.attn = HeadFusionAttentionV2(
+                dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        else:
+            self.attn = Attention(
+                dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        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)
+
+    def forward(self, x):
+        feature = self.attn(self.norm1(x))
+        x = x + self.drop_path(feature)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+

CAGE_expression_inference-apvit/apvit_mmcls/models/backbones/modules/vit_pooling.py

@@ -0,0 +1,89 @@
+import torch
+from torch import nn
+from .vit import Mlp
+from mmcls.models.utils import top_pool
+from mmcls.models.vit.layers import DropPath
+
+
+class PoolingAttention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., 
+                 pool_config=None):
+        super().__init__()
+        self.num_heads = num_heads
+        self.dim = dim
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        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)
+        self.pool_config = pool_config
+
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale   # [B, head_num, token_num, token_num]
+
+        if self.pool_config:
+            attn_method = self.pool_config.get('attn_method')
+            if attn_method == 'SUM_ABS_1':
+                attn_weight = attn[:, :, 0, :].transpose(-1, -2)    # [B, token_num, head_num]
+                attn_weight = torch.sum(torch.abs(attn_weight), dim=-1).unsqueeze(-1)
+            elif attn_method == 'SUM':
+                attn_weight = attn[:, :, 0, :].transpose(-1, -2)    # [B, token_num, head_num]
+                attn_weight = torch.sum(attn_weight, dim=-1).unsqueeze(-1)
+            elif attn_method == 'MAX':
+                attn_weight = attn[:, :, 0, :].transpose(-1, -2)
+                attn_weight = torch.max(attn_weight, dim=-1)[0].unsqueeze(-1)
+            else:
+                raise ValueError('Invalid attn_method: %s' % attn_method)
+
+            # attn_weight = torch.rand(attn_weight.shape, device=attn_weight.device)
+            keep_index = top_pool(attn_weight, dim=self.dim, **self.pool_config)
+        else:
+            keep_index = None
+
+        attn = attn.softmax(dim=-1)
+        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, keep_index
+
+
+class PoolingBlock(nn.Module):
+
+    def __init__(self, dim=0, num_heads=0, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, pool_config=None, **kwargs):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = PoolingAttention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, 
+            pool_config=pool_config)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        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)
+
+    def forward(self, x):
+        feature, keep_index = self.attn(self.norm1(x))
+        x = x + self.drop_path(feature)
+        if keep_index is not None:
+            if len(keep_index) != x.shape[1]:
+                x = x.gather(dim=1, index=keep_index)
+                # pooled_x = []
+                # for i in range(keep_index.shape[0]):
+                #     pooled_x.append(x[i, keep_index[i, :, 0]])
+                # x = torch.stack(pooled_x)
+                # assert torch.all(torch.eq(quick_x, x))
+
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x