Initial commit

classification/configs/intern_vit_6b_1k_224_test_imagenetv2.yaml

+DATA:
+  IMG_ON_MEMORY: False
+  BATCH_SIZE: 128
+  DATASET: 'imagenetv2'
+  TRANSFORM: 'build_transform_for_linear_probe'
+  DATA_PATH: './data/imagenetv2'
+MODEL:
+  TYPE: intern_vit_6b
+  DROP_PATH_RATE: 0.0
+  INTERN_VIT_6B:
+    FREEZE_VIT: True
+    PATCH_SIZE: 14
+    PRETRAIN_SIZE: 224
+    QKV_BIAS: False
+    EMBED_DIM: 3200
+    NUM_HEADS: 25
+    MLP_RATIO: 4
+    INIT_VALUES: 0.1
+    QK_NORMALIZATION: True
+    DEPTH: 48
+    USE_FLASH_ATTN: True
+    PRETRAINED: "./pretrained/intern_vit_6b_224px.pth"
+    CLS_TARGET: 'cls_patch_concat'
+TRAIN:
+  EMA:
+    ENABLE: False
+    DECAY: 0.998
+  EPOCHS: 10
+  WARMUP_EPOCHS: 1
+  WEIGHT_DECAY: 0.0
+  BASE_LR: 0.1 # 512
+  WARMUP_LR: .0
+  MIN_LR: .0
+  LR_LAYER_DECAY: false
+  OPTIMIZER:
+    NAME: 'sgd'

classification/dataset/__init__.py

+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from .build import build_loader, build_loader2

classification/dataset/build.py

+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import os
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from timm.data import Mixup, create_transform
+from torchvision import transforms
+from torchvision.datasets import ImageFolder
+
+from .cached_image_folder import ImageCephDataset
+from .samplers import NodeDistributedSampler, SubsetRandomSampler
+
+try:
+    from torchvision.transforms import InterpolationMode
+
+    def _pil_interp(method):
+        if method == 'bicubic':
+            return InterpolationMode.BICUBIC
+        elif method == 'lanczos':
+            return InterpolationMode.LANCZOS
+        elif method == 'hamming':
+            return InterpolationMode.HAMMING
+        else:
+            return InterpolationMode.BILINEAR
+except:
+    from timm.data.transforms import _pil_interp
+
+
+class TTA(torch.nn.Module):
+
+    def __init__(self, size, scales=[1.0, 1.05, 1.1]):
+        super().__init__()
+        self.size = size
+        self.scales = scales
+
+    def forward(self, img):
+        out = []
+        cc = transforms.CenterCrop(self.size)
+        for scale in self.scales:
+            size_ = int(scale * self.size)
+            rs = transforms.Resize(size_, interpolation=_pil_interp('bicubic'))
+            img_ = rs(img)
+            img_ = cc(img_)
+            out.append(img_)
+
+        return out
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(size={self.size}, scale={self.scales})'
+
+
+def build_loader(config):
+    config.defrost()
+    dataset_train, config.MODEL.NUM_CLASSES = build_dataset('train', config=config)
+    config.freeze()
+    print(f'local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}'
+          'successfully build train dataset')
+
+    dataset_val, _ = build_dataset('val', config=config)
+    print(f'local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}'
+          'successfully build val dataset')
+
+    dataset_test, _ = build_dataset('test', config=config)
+    print(f'local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}'
+          'successfully build test dataset')
+
+    num_tasks = dist.get_world_size()
+    global_rank = dist.get_rank()
+
+    if dataset_train is not None:
+        if config.DATA.IMG_ON_MEMORY:
+            sampler_train = NodeDistributedSampler(dataset_train)
+        else:
+            if config.DATA.ZIP_MODE and config.DATA.CACHE_MODE == 'part':
+                indices = np.arange(dist.get_rank(), len(dataset_train), dist.get_world_size())
+                sampler_train = SubsetRandomSampler(indices)
+            else:
+                sampler_train = torch.utils.data.DistributedSampler(
+                    dataset_train,
+                    num_replicas=num_tasks,
+                    rank=global_rank,
+                    shuffle=True)
+
+    if dataset_val is not None:
+        if config.TEST.SEQUENTIAL:
+            sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+        else:
+            sampler_val = torch.utils.data.distributed.DistributedSampler(dataset_val, shuffle=False)
+
+    if dataset_test is not None:
+        if config.TEST.SEQUENTIAL:
+            sampler_test = torch.utils.data.SequentialSampler(dataset_test)
+        else:
+            sampler_test = torch.utils.data.distributed.DistributedSampler(dataset_test, shuffle=False)
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train,
+        sampler=sampler_train,
+        batch_size=config.DATA.BATCH_SIZE,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=True,
+        persistent_workers=True) if dataset_train is not None else None
+
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val,
+        sampler=sampler_val,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_val is not None else None
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test,
+        sampler=sampler_test,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_test is not None else None
+
+    # setup mixup / cutmix
+    mixup_fn = None
+    mixup_active = config.AUG.MIXUP > 0 or config.AUG.CUTMIX > 0. or config.AUG.CUTMIX_MINMAX is not None
+    if mixup_active:
+        mixup_fn = Mixup(mixup_alpha=config.AUG.MIXUP,
+                         cutmix_alpha=config.AUG.CUTMIX,
+                         cutmix_minmax=config.AUG.CUTMIX_MINMAX,
+                         prob=config.AUG.MIXUP_PROB,
+                         switch_prob=config.AUG.MIXUP_SWITCH_PROB,
+                         mode=config.AUG.MIXUP_MODE,
+                         label_smoothing=config.MODEL.LABEL_SMOOTHING,
+                         num_classes=config.MODEL.NUM_CLASSES)
+
+    return dataset_train, dataset_val, dataset_test, data_loader_train, \
+           data_loader_val, data_loader_test, mixup_fn
+
+
+def build_loader2(config):
+    config.defrost()
+    dataset_train, config.MODEL.NUM_CLASSES = build_dataset('train', config=config)
+    config.freeze()
+    dataset_val, _ = build_dataset('val', config=config)
+    dataset_test, _ = build_dataset('test', config=config)
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train,
+        shuffle=True,
+        batch_size=config.DATA.BATCH_SIZE,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=True,
+        persistent_workers=True) if dataset_train is not None else None
+
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_val is not None else None
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_test is not None else None
+
+    # setup mixup / cutmix
+    mixup_fn = None
+    mixup_active = config.AUG.MIXUP > 0 or config.AUG.CUTMIX > 0. or config.AUG.CUTMIX_MINMAX is not None
+    if mixup_active:
+        mixup_fn = Mixup(mixup_alpha=config.AUG.MIXUP,
+                         cutmix_alpha=config.AUG.CUTMIX,
+                         cutmix_minmax=config.AUG.CUTMIX_MINMAX,
+                         prob=config.AUG.MIXUP_PROB,
+                         switch_prob=config.AUG.MIXUP_SWITCH_PROB,
+                         mode=config.AUG.MIXUP_MODE,
+                         label_smoothing=config.MODEL.LABEL_SMOOTHING,
+                         num_classes=config.MODEL.NUM_CLASSES)
+
+    return dataset_train, dataset_val, dataset_test, data_loader_train, \
+           data_loader_val, data_loader_test, mixup_fn
+
+
+def build_dataset(split, config):
+    if config.DATA.TRANSFORM == 'build_transform':
+        transform = build_transform(split == 'train', config)
+    elif config.DATA.TRANSFORM == 'build_transform_for_linear_probe':
+        transform = build_transform_for_linear_probe(split == 'train', config)
+    else:
+        raise NotImplementedError
+    print(split, transform)
+    dataset = None
+    nb_classes = None
+    prefix = split
+    if config.DATA.DATASET == 'imagenet' or config.DATA.DATASET == 'imagenet-real':
+        if prefix == 'train' and not config.EVAL_MODE:
+            root = os.path.join(config.DATA.DATA_PATH, 'train')
+            dataset = ImageCephDataset(root, 'train',
+                                       transform=transform,
+                                       on_memory=config.DATA.IMG_ON_MEMORY)
+        elif prefix == 'val':
+            root = os.path.join(config.DATA.DATA_PATH, 'val')
+            dataset = ImageCephDataset(root, 'val', transform=transform)
+        nb_classes = 1000
+    elif config.DATA.DATASET == 'imagenet22K':
+        if prefix == 'train':
+            if not config.EVAL_MODE:
+                root = config.DATA.DATA_PATH
+                dataset = ImageCephDataset(root, 'train',
+                                           transform=transform,
+                                           on_memory=config.DATA.IMG_ON_MEMORY)
+            nb_classes = 21841
+        elif prefix == 'val':
+            root = os.path.join(config.DATA.DATA_PATH, 'val')
+            dataset = ImageCephDataset(root, 'val', transform=transform)
+            nb_classes = 1000
+    elif config.DATA.DATASET == 'imagenetv2':
+        from .imagenetv2 import ImageNetV2Dataset
+        if prefix == 'train' and not config.EVAL_MODE:
+            print(f'Only test split available for {config.DATA.DATASET}')
+        else:
+            dataset = ImageNetV2Dataset(variant='matched-frequency',
+                                        transform=transform,
+                                        location=config.DATA.DATA_PATH)
+            nb_classes = 1000
+    elif config.DATA.DATASET == 'imagenet_sketch':
+        if prefix == 'train' and not config.EVAL_MODE:
+            print(f'Only test split available for {config.DATA.DATASET}')
+        else:
+            dataset = ImageFolder(root=config.DATA.DATA_PATH, transform=transform)
+            nb_classes = 1000
+    elif config.DATA.DATASET == 'imagenet_a':
+        if prefix == 'train' and not config.EVAL_MODE:
+            print(f'Only test split available for {config.DATA.DATASET}')
+        else:
+            dataset = ImageFolder(root=config.DATA.DATA_PATH, transform=transform)
+            nb_classes = 1000  # actual number of classes is 200
+    elif config.DATA.DATASET == 'imagenet_r':
+        if prefix == 'train' and not config.EVAL_MODE:
+            print(f'Only test split available for {config.DATA.DATASET}')
+        else:
+            dataset = ImageFolder(root=config.DATA.DATA_PATH, transform=transform)
+            nb_classes = 1000  # actual number of classes is 200
+    else:
+        raise NotImplementedError(
+            f'build_dataset does support {config.DATA.DATASET}')
+
+    return dataset, nb_classes
+
+
+def build_transform_for_linear_probe(is_train, config):
+    # linear probe: weak augmentation
+    if is_train:
+        transform = transforms.Compose([
+            transforms.RandomResizedCrop(
+                config.DATA.IMG_SIZE, interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=config.AUG.MEAN, std=config.AUG.STD)
+        ])
+    else:
+        transform = transforms.Compose([
+            transforms.Resize(
+                config.DATA.IMG_SIZE, interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.CenterCrop(config.DATA.IMG_SIZE),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=config.AUG.MEAN, std=config.AUG.STD)
+        ])
+    return transform
+
+
+def build_transform(is_train, config):
+    resize_im = config.DATA.IMG_SIZE > 32
+    if is_train:
+        # this should always dispatch to transforms_imagenet_train
+        transform = create_transform(
+            input_size=config.DATA.IMG_SIZE,
+            is_training=True,
+            color_jitter=config.AUG.COLOR_JITTER
+            if config.AUG.COLOR_JITTER > 0 else None,
+            auto_augment=config.AUG.AUTO_AUGMENT
+            if config.AUG.AUTO_AUGMENT != 'none' else None,
+            re_prob=config.AUG.REPROB,
+            re_mode=config.AUG.REMODE,
+            re_count=config.AUG.RECOUNT,
+            interpolation=config.DATA.INTERPOLATION,
+        )
+        if not resize_im:
+            # replace RandomResizedCropAndInterpolation with
+            # RandomCrop
+            transform.transforms[0] = transforms.RandomCrop(config.DATA.IMG_SIZE, padding=4)
+
+        return transform
+
+    t = []
+    if resize_im:
+        if config.TEST.CROP:
+            size = int(1.0 * config.DATA.IMG_SIZE)
+            t.append(
+                transforms.Resize(size, interpolation=_pil_interp(config.DATA.INTERPOLATION)),
+                # to maintain same ratio w.r.t. 224 images
+            )
+            t.append(transforms.CenterCrop(config.DATA.IMG_SIZE))
+        elif config.AUG.RANDOM_RESIZED_CROP:
+            t.append(
+                transforms.RandomResizedCrop(
+                    (config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
+                    interpolation=_pil_interp(config.DATA.INTERPOLATION)))
+        else:
+            t.append(
+                transforms.Resize(
+                    (config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
+                    interpolation=_pil_interp(config.DATA.INTERPOLATION)))
+    t.append(transforms.ToTensor())
+    t.append(transforms.Normalize(config.AUG.MEAN, config.AUG.STD))
+
+    return transforms.Compose(t)

classification/dataset/cached_image_folder.py

+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import io
+import json
+import logging
+import math
+import os
+import os.path as osp
+import re
+import time
+from abc import abstractmethod
+
+import mmcv
+import torch
+import torch.distributed as dist
+import torch.utils.data as data
+from mmcv.fileio import FileClient
+from PIL import Image
+from tqdm import tqdm, trange
+
+from .zipreader import ZipReader, is_zip_path
+
+_logger = logging.getLogger(__name__)
+
+_ERROR_RETRY = 50
+
+
+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_classes(dir):
+    classes = [
+        d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))
+    ]
+    classes.sort()
+    class_to_idx = {classes[i]: i for i in range(len(classes))}
+    return classes, class_to_idx
+
+
+def make_dataset(dir, class_to_idx, extensions):
+    images = []
+    dir = os.path.expanduser(dir)
+    for target in sorted(os.listdir(dir)):
+        d = os.path.join(dir, target)
+        if not os.path.isdir(d):
+            continue
+        for root, _, fnames in sorted(os.walk(d)):
+            for fname in sorted(fnames):
+                if has_file_allowed_extension(fname, extensions):
+                    path = os.path.join(root, fname)
+                    item = (path, class_to_idx[target])
+                    images.append(item)
+
+    return images
+
+
+def make_dataset_with_ann(ann_file, img_prefix, extensions):
+    images = []
+    with open(ann_file, 'r') as f:
+        contents = f.readlines()
+        for line_str in contents:
+            path_contents = [c for c in line_str.split('\t')]
+            im_file_name = path_contents[0]
+            class_index = int(path_contents[1])
+            assert str.lower(os.path.splitext(im_file_name)[-1]) in extensions
+            item = (os.path.join(img_prefix, im_file_name), class_index)
+            images.append(item)
+
+    return images
+
+
+class DatasetFolder(data.Dataset):
+    """A generic data loader where the samples are arranged in this way: ::
+
+    root/class_x/xxx.ext
+        root/class_x/xxy.ext
+        root/class_x/xxz.ext
+        root/class_y/123.ext
+        root/class_y/nsdf3.ext
+        root/class_y/asd932_.ext
+    Args:
+        root (string): Root directory path.
+        loader (callable): A function to load a sample given its path.
+        extensions (list[string]): A list of allowed extensions.
+        transform (callable, optional): A function/transform that takes in
+            a sample and returns a transformed version.
+            E.g, ``transforms.RandomCrop`` for images.
+        target_transform (callable, optional): A function/transform that takes
+            in the target and transforms it.
+     Attributes:
+        samples (list): List of (sample path, class_index) tuples
+    """
+
+    def __init__(self,
+                 root,
+                 loader,
+                 extensions,
+                 ann_file='',
+                 img_prefix='',
+                 transform=None,
+                 target_transform=None,
+                 cache_mode='no'):
+        # image folder mode
+        if ann_file == '':
+            _, class_to_idx = find_classes(root)
+            samples = make_dataset(root, class_to_idx, extensions)
+        # zip mode
+        else:
+            samples = make_dataset_with_ann(os.path.join(root, ann_file),
+                                            os.path.join(root, img_prefix),
+                                            extensions)
+
+        if len(samples) == 0:
+            raise (RuntimeError('Found 0 files in subfolders of: ' + root +
+                                '\n' + 'Supported extensions are: ' +
+                                ','.join(extensions)))
+
+        self.root = root
+        self.loader = loader
+        self.extensions = extensions
+
+        self.samples = samples
+        self.labels = [y_1k for _, y_1k in samples]
+        self.classes = list(set(self.labels))
+
+        self.transform = transform
+        self.target_transform = target_transform
+
+        self.cache_mode = cache_mode
+        if self.cache_mode != 'no':
+            self.init_cache()
+
+    def init_cache(self):
+        assert self.cache_mode in ['part', 'full']
+        n_sample = len(self.samples)
+        global_rank = dist.get_rank()
+        world_size = dist.get_world_size()
+
+        samples_bytes = [None for _ in range(n_sample)]
+        start_time = time.time()
+        for index in range(n_sample):
+            if index % (n_sample // 10) == 0:
+                t = time.time() - start_time
+                print(
+                    f'global_rank {dist.get_rank()} cached {index}/{n_sample} takes {t:.2f}s per block'
+                )
+                start_time = time.time()
+            path, target = self.samples[index]
+            if self.cache_mode == 'full':
+                samples_bytes[index] = (ZipReader.read(path), target)
+            elif self.cache_mode == 'part' and index % world_size == global_rank:
+                samples_bytes[index] = (ZipReader.read(path), target)
+            else:
+                samples_bytes[index] = (path, target)
+        self.samples = samples_bytes
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (sample, target) where target is class_index of the target class.
+        """
+        path, target = self.samples[index]
+        sample = self.loader(path)
+        if self.transform is not None:
+            sample = self.transform(sample)
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return sample, target
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __repr__(self):
+        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
+        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
+        fmt_str += '    Root Location: {}\n'.format(self.root)
+        tmp = '    Transforms (if any): '
+        fmt_str += '{0}{1}\n'.format(
+            tmp,
+            self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
+        tmp = '    Target Transforms (if any): '
+        fmt_str += '{0}{1}'.format(
+            tmp,
+            self.target_transform.__repr__().replace('\n',
+                                                     '\n' + ' ' * len(tmp)))
+
+        return fmt_str
+
+
+IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']
+
+
+def pil_loader(path):
+    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
+    if isinstance(path, bytes):
+        img = Image.open(io.BytesIO(path))
+    elif is_zip_path(path):
+        data = ZipReader.read(path)
+        img = Image.open(io.BytesIO(data))
+    else:
+        with open(path, 'rb') as f:
+            img = Image.open(f)
+            return img.convert('RGB')
+
+    return img.convert('RGB')
+
+
+def accimage_loader(path):
+    import accimage
+    try:
+        return accimage.Image(path)
+    except IOError:
+        # Potentially a decoding problem, fall back to PIL.Image
+        return pil_loader(path)
+
+
+def default_img_loader(path):
+    from torchvision import get_image_backend
+    if get_image_backend() == 'accimage':
+        return accimage_loader(path)
+    else:
+        return pil_loader(path)
+
+
+class CachedImageFolder(DatasetFolder):
+    """A generic data loader where the images are arranged in this way: ::
+
+    root/dog/xxx.png
+        root/dog/xxy.png
+        root/dog/xxz.png
+        root/cat/123.png
+        root/cat/nsdf3.png
+        root/cat/asd932_.png
+    Args:
+        root (string): Root directory path.
+        transform (callable, optional): A function/transform that  takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.RandomCrop``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+        loader (callable, optional): A function to load an image given its path.
+     Attributes:
+        imgs (list): List of (image path, class_index) tuples
+    """
+
+    def __init__(self,
+                 root,
+                 ann_file='',
+                 img_prefix='',
+                 transform=None,
+                 target_transform=None,
+                 loader=default_img_loader,
+                 cache_mode='no'):
+        super(CachedImageFolder,
+              self).__init__(root,
+                             loader,
+                             IMG_EXTENSIONS,
+                             ann_file=ann_file,
+                             img_prefix=img_prefix,
+                             transform=transform,
+                             target_transform=target_transform,
+                             cache_mode=cache_mode)
+        self.imgs = self.samples
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (image, target) where target is class_index of the target class.
+        """
+        path, target = self.samples[index]
+        image = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(image)
+        else:
+            img = image
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
+
+
+class ImageCephDataset(data.Dataset):
+
+    def __init__(self,
+                 root,
+                 split,
+                 parser=None,
+                 transform=None,
+                 target_transform=None,
+                 on_memory=False):
+        if '22k' in root:
+            # Imagenet 22k
+            annotation_root = 'meta_data/'
+        else:
+            # Imagenet
+            annotation_root = 'meta_data/'
+        if parser is None or isinstance(parser, str):
+            parser = ParserCephImage(root=root,
+                                     split=split,
+                                     annotation_root=annotation_root,
+                                     on_memory=on_memory)
+        self.parser = parser
+        self.transform = transform
+        self.target_transform = target_transform
+        self._consecutive_errors = 0
+
+    def __getitem__(self, index):
+        img, target = self.parser[index]
+        self._consecutive_errors = 0
+        if self.transform is not None:
+            img = self.transform(img)
+        if target is None:
+            target = -1
+        elif self.target_transform is not None:
+            target = self.target_transform(target)
+        return img, target
+
+    def __len__(self):
+        return len(self.parser)
+
+    def filename(self, index, basename=False, absolute=False):
+        return self.parser.filename(index, basename, absolute)
+
+    def filenames(self, basename=False, absolute=False):
+        return self.parser.filenames(basename, absolute)
+
+
+class Parser:
+
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def _filename(self, index, basename=False, absolute=False):
+        pass
+
+    def filename(self, index, basename=False, absolute=False):
+        return self._filename(index, basename=basename, absolute=absolute)
+
+    def filenames(self, basename=False, absolute=False):
+        return [
+            self._filename(index, basename=basename, absolute=absolute)
+            for index in range(len(self))
+        ]
+
+
+class ParserCephImage(Parser):
+
+    def __init__(self,
+                 root,
+                 split,
+                 annotation_root,
+                 on_memory=False,
+                 **kwargs):
+        super().__init__()
+
+        self.file_client = None
+        self.kwargs = kwargs
+
+        self.root = root  # dataset:s3://imagenet22k
+        if '22k' in root:
+            self.io_backend = 'petrel'
+            with open(osp.join(annotation_root, '22k_class_to_idx.json'),
+                      'r') as f:
+                self.class_to_idx = json.loads(f.read())
+            with open(osp.join(annotation_root, '22k_label.txt'), 'r') as f:
+                self.samples = f.read().splitlines()
+        else:
+            self.io_backend = 'disk'
+            self.class_to_idx = None
+            with open(osp.join(annotation_root, f'{split}.txt'), 'r') as f:
+                self.samples = f.read().splitlines()
+        local_rank = None
+        local_size = None
+        self._consecutive_errors = 0
+        self.on_memory = on_memory
+        if on_memory:
+            self.holder = {}
+            if local_rank is None:
+                local_rank = int(os.environ.get('LOCAL_RANK', 0))
+            if local_size is None:
+                local_size = int(os.environ.get('LOCAL_SIZE', 1))
+            self.local_rank = local_rank
+            self.local_size = local_size
+            self.rank = int(os.environ['RANK'])
+            self.world_size = int(os.environ['WORLD_SIZE'])
+            self.num_replicas = int(os.environ['WORLD_SIZE'])
+            self.num_parts = local_size
+            self.num_samples = int(
+                math.ceil(len(self.samples) * 1.0 / self.num_replicas))
+            self.total_size = self.num_samples * self.num_replicas
+            self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
+            self.load_onto_memory_v2()
+
+    def load_onto_memory(self):
+        print('Loading images onto memory...', self.local_rank,
+              self.local_size)
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend, **self.kwargs)
+        for index in trange(len(self.samples)):
+            if index % self.local_size != self.local_rank:
+                continue
+            path, _ = self.samples[index].split(' ')
+            path = osp.join(self.root, path)
+            img_bytes = self.file_client.get(path)
+            self.holder[path] = img_bytes
+
+        print('Loading complete!')
+
+    def load_onto_memory_v2(self):
+        # print("Loading images onto memory...", self.local_rank, self.local_size)
+        t = torch.Generator()
+        t.manual_seed(0)
+        indices = torch.randperm(len(self.samples), generator=t).tolist()
+        # indices = range(len(self.samples))
+        indices = [i for i in indices if i % self.num_parts == self.local_rank]
+        # add extra samples to make it evenly divisible
+        indices += indices[:(self.total_size_parts - len(indices))]
+        assert len(indices) == self.total_size_parts
+
+        # subsample
+        indices = indices[self.rank // self.num_parts:self.
+                          total_size_parts:self.num_replicas // self.num_parts]
+        assert len(indices) == self.num_samples
+
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend, **self.kwargs)
+        for index in tqdm(indices):
+            if index % self.local_size != self.local_rank:
+                continue
+            path, _ = self.samples[index].split(' ')
+            path = osp.join(self.root, path)
+            img_bytes = self.file_client.get(path)
+
+            self.holder[path] = img_bytes
+
+        print('Loading complete!')
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend, **self.kwargs)
+
+        filepath, target = self.samples[index].split(' ')
+        filepath = osp.join(self.root, filepath)
+
+        try:
+            if self.on_memory:
+                img_bytes = self.holder[filepath]
+            else:
+                # pass
+                img_bytes = self.file_client.get(filepath)
+            img = mmcv.imfrombytes(img_bytes)[:, :, ::-1]
+        except Exception as e:
+            _logger.warning(
+                f'Skipped sample (index {index}, file {filepath}). {str(e)}')
+            self._consecutive_errors += 1
+            if self._consecutive_errors < _ERROR_RETRY:
+                return self.__getitem__((index + 1) % len(self))
+            else:
+                raise e
+        self._consecutive_errors = 0
+
+        img = Image.fromarray(img)
+        try:
+            if self.class_to_idx is not None:
+                target = self.class_to_idx[target]
+            else:
+                target = int(target)
+        except:
+            print(filepath, target)
+            exit()
+
+        return img, target
+
+    def __len__(self):
+        return len(self.samples)
+
+    def _filename(self, index, basename=False, absolute=False):
+        filename, _ = self.samples[index].split(' ')
+        filename = osp.join(self.root, filename)
+
+        return filename
+
+
+def get_temporal_info(date, miss_hour=False):
+    try:
+        if date:
+            if miss_hour:
+                pattern = re.compile(r'(\d*)-(\d*)-(\d*)', re.I)
+            else:
+                pattern = re.compile(r'(\d*)-(\d*)-(\d*) (\d*):(\d*):(\d*)',
+                                     re.I)
+            m = pattern.match(date.strip())
+
+            if m:
+                year = int(m.group(1))
+                month = int(m.group(2))
+                day = int(m.group(3))
+                x_month = math.sin(2 * math.pi * month / 12)
+                y_month = math.cos(2 * math.pi * month / 12)
+                if miss_hour:
+                    x_hour = 0
+                    y_hour = 0
+                else:
+                    hour = int(m.group(4))
+                    x_hour = math.sin(2 * math.pi * hour / 24)
+                    y_hour = math.cos(2 * math.pi * hour / 24)
+                return [x_month, y_month, x_hour, y_hour]
+            else:
+                return [0, 0, 0, 0]
+        else:
+            return [0, 0, 0, 0]
+    except:
+        return [0, 0, 0, 0]
+
+
+def get_spatial_info(latitude, longitude):
+    if latitude and longitude:
+        latitude = math.radians(latitude)
+        longitude = math.radians(longitude)
+        x = math.cos(latitude) * math.cos(longitude)
+        y = math.cos(latitude) * math.sin(longitude)
+        z = math.sin(latitude)
+        return [x, y, z]
+    else:
+        return [0, 0, 0]

classification/dataset/imagenet_a_r_indices.py

+"""Code from https://github.com/baaivision/EVA/blob/master/EVA-02/asuka/imagenet_a_r_indices.py
+Thanks to the authors of EVA."""
+
+all_wnids = [
+    'n01440764', 'n01443537', 'n01484850', 'n01491361', 'n01494475',
+    'n01496331', 'n01498041', 'n01514668', 'n01514859', 'n01518878',
+    'n01530575', 'n01531178', 'n01532829', 'n01534433', 'n01537544',
+    'n01558993', 'n01560419', 'n01580077', 'n01582220', 'n01592084',
+    'n01601694', 'n01608432', 'n01614925', 'n01616318', 'n01622779',
+    'n01629819', 'n01630670', 'n01631663', 'n01632458', 'n01632777',
+    'n01641577', 'n01644373', 'n01644900', 'n01664065', 'n01665541',
+    'n01667114', 'n01667778', 'n01669191', 'n01675722', 'n01677366',
+    'n01682714', 'n01685808', 'n01687978', 'n01688243', 'n01689811',
+    'n01692333', 'n01693334', 'n01694178', 'n01695060', 'n01697457',
+    'n01698640', 'n01704323', 'n01728572', 'n01728920', 'n01729322',
+    'n01729977', 'n01734418', 'n01735189', 'n01737021', 'n01739381',
+    'n01740131', 'n01742172', 'n01744401', 'n01748264', 'n01749939',
+    'n01751748', 'n01753488', 'n01755581', 'n01756291', 'n01768244',
+    'n01770081', 'n01770393', 'n01773157', 'n01773549', 'n01773797',
+    'n01774384', 'n01774750', 'n01775062', 'n01776313', 'n01784675',
+    'n01795545', 'n01796340', 'n01797886', 'n01798484', 'n01806143',
+    'n01806567', 'n01807496', 'n01817953', 'n01818515', 'n01819313',
+    'n01820546', 'n01824575', 'n01828970', 'n01829413', 'n01833805',
+    'n01843065', 'n01843383', 'n01847000', 'n01855032', 'n01855672',
+    'n01860187', 'n01871265', 'n01872401', 'n01873310', 'n01877812',
+    'n01882714', 'n01883070', 'n01910747', 'n01914609', 'n01917289',
+    'n01924916', 'n01930112', 'n01943899', 'n01944390', 'n01945685',
+    'n01950731', 'n01955084', 'n01968897', 'n01978287', 'n01978455',
+    'n01980166', 'n01981276', 'n01983481', 'n01984695', 'n01985128',
+    'n01986214', 'n01990800', 'n02002556', 'n02002724', 'n02006656',
+    'n02007558', 'n02009229', 'n02009912', 'n02011460', 'n02012849',
+    'n02013706', 'n02017213', 'n02018207', 'n02018795', 'n02025239',
+    'n02027492', 'n02028035', 'n02033041', 'n02037110', 'n02051845',
+    'n02056570', 'n02058221', 'n02066245', 'n02071294', 'n02074367',
+    'n02077923', 'n02085620', 'n02085782', 'n02085936', 'n02086079',
+    'n02086240', 'n02086646', 'n02086910', 'n02087046', 'n02087394',
+    'n02088094', 'n02088238', 'n02088364', 'n02088466', 'n02088632',
+    'n02089078', 'n02089867', 'n02089973', 'n02090379', 'n02090622',
+    'n02090721', 'n02091032', 'n02091134', 'n02091244', 'n02091467',
+    'n02091635', 'n02091831', 'n02092002', 'n02092339', 'n02093256',
+    'n02093428', 'n02093647', 'n02093754', 'n02093859', 'n02093991',
+    'n02094114', 'n02094258', 'n02094433', 'n02095314', 'n02095570',
+    'n02095889', 'n02096051', 'n02096177', 'n02096294', 'n02096437',
+    'n02096585', 'n02097047', 'n02097130', 'n02097209', 'n02097298',
+    'n02097474', 'n02097658', 'n02098105', 'n02098286', 'n02098413',
+    'n02099267', 'n02099429', 'n02099601', 'n02099712', 'n02099849',
+    'n02100236', 'n02100583', 'n02100735', 'n02100877', 'n02101006',
+    'n02101388', 'n02101556', 'n02102040', 'n02102177', 'n02102318',
+    'n02102480', 'n02102973', 'n02104029', 'n02104365', 'n02105056',
+    'n02105162', 'n02105251', 'n02105412', 'n02105505', 'n02105641',
+    'n02105855', 'n02106030', 'n02106166', 'n02106382', 'n02106550',
+    'n02106662', 'n02107142', 'n02107312', 'n02107574', 'n02107683',
+    'n02107908', 'n02108000', 'n02108089', 'n02108422', 'n02108551',
+    'n02108915', 'n02109047', 'n02109525', 'n02109961', 'n02110063',
+    'n02110185', 'n02110341', 'n02110627', 'n02110806', 'n02110958',
+    'n02111129', 'n02111277', 'n02111500', 'n02111889', 'n02112018',
+    'n02112137', 'n02112350', 'n02112706', 'n02113023', 'n02113186',
+    'n02113624', 'n02113712', 'n02113799', 'n02113978', 'n02114367',
+    'n02114548', 'n02114712', 'n02114855', 'n02115641', 'n02115913',
+    'n02116738', 'n02117135', 'n02119022', 'n02119789', 'n02120079',
+    'n02120505', 'n02123045', 'n02123159', 'n02123394', 'n02123597',
+    'n02124075', 'n02125311', 'n02127052', 'n02128385', 'n02128757',
+    'n02128925', 'n02129165', 'n02129604', 'n02130308', 'n02132136',
+    'n02133161', 'n02134084', 'n02134418', 'n02137549', 'n02138441',
+    'n02165105', 'n02165456', 'n02167151', 'n02168699', 'n02169497',
+    'n02172182', 'n02174001', 'n02177972', 'n02190166', 'n02206856',
+    'n02219486', 'n02226429', 'n02229544', 'n02231487', 'n02233338',
+    'n02236044', 'n02256656', 'n02259212', 'n02264363', 'n02268443',
+    'n02268853', 'n02276258', 'n02277742', 'n02279972', 'n02280649',
+    'n02281406', 'n02281787', 'n02317335', 'n02319095', 'n02321529',
+    'n02325366', 'n02326432', 'n02328150', 'n02342885', 'n02346627',
+    'n02356798', 'n02361337', 'n02363005', 'n02364673', 'n02389026',
+    'n02391049', 'n02395406', 'n02396427', 'n02397096', 'n02398521',
+    'n02403003', 'n02408429', 'n02410509', 'n02412080', 'n02415577',
+    'n02417914', 'n02422106', 'n02422699', 'n02423022', 'n02437312',
+    'n02437616', 'n02441942', 'n02442845', 'n02443114', 'n02443484',
+    'n02444819', 'n02445715', 'n02447366', 'n02454379', 'n02457408',
+    'n02480495', 'n02480855', 'n02481823', 'n02483362', 'n02483708',
+    'n02484975', 'n02486261', 'n02486410', 'n02487347', 'n02488291',
+    'n02488702', 'n02489166', 'n02490219', 'n02492035', 'n02492660',
+    'n02493509', 'n02493793', 'n02494079', 'n02497673', 'n02500267',
+    'n02504013', 'n02504458', 'n02509815', 'n02510455', 'n02514041',
+    'n02526121', 'n02536864', 'n02606052', 'n02607072', 'n02640242',
+    'n02641379', 'n02643566', 'n02655020', 'n02666196', 'n02667093',
+    'n02669723', 'n02672831', 'n02676566', 'n02687172', 'n02690373',
+    'n02692877', 'n02699494', 'n02701002', 'n02704792', 'n02708093',
+    'n02727426', 'n02730930', 'n02747177', 'n02749479', 'n02769748',
+    'n02776631', 'n02777292', 'n02782093', 'n02783161', 'n02786058',
+    'n02787622', 'n02788148', 'n02790996', 'n02791124', 'n02791270',
+    'n02793495', 'n02794156', 'n02795169', 'n02797295', 'n02799071',
+    'n02802426', 'n02804414', 'n02804610', 'n02807133', 'n02808304',
+    'n02808440', 'n02814533', 'n02814860', 'n02815834', 'n02817516',
+    'n02823428', 'n02823750', 'n02825657', 'n02834397', 'n02835271',
+    'n02837789', 'n02840245', 'n02841315', 'n02843684', 'n02859443',
+    'n02860847', 'n02865351', 'n02869837', 'n02870880', 'n02871525',
+    'n02877765', 'n02879718', 'n02883205', 'n02892201', 'n02892767',
+    'n02894605', 'n02895154', 'n02906734', 'n02909870', 'n02910353',
+    'n02916936', 'n02917067', 'n02927161', 'n02930766', 'n02939185',
+    'n02948072', 'n02950826', 'n02951358', 'n02951585', 'n02963159',
+    'n02965783', 'n02966193', 'n02966687', 'n02971356', 'n02974003',
+    'n02977058', 'n02978881', 'n02979186', 'n02980441', 'n02981792',
+    'n02988304', 'n02992211', 'n02992529', 'n02999410', 'n03000134',
+    'n03000247', 'n03000684', 'n03014705', 'n03016953', 'n03017168',
+    'n03018349', 'n03026506', 'n03028079', 'n03032252', 'n03041632',
+    'n03042490', 'n03045698', 'n03047690', 'n03062245', 'n03063599',
+    'n03063689', 'n03065424', 'n03075370', 'n03085013', 'n03089624',
+    'n03095699', 'n03100240', 'n03109150', 'n03110669', 'n03124043',
+    'n03124170', 'n03125729', 'n03126707', 'n03127747', 'n03127925',
+    'n03131574', 'n03133878', 'n03134739', 'n03141823', 'n03146219',
+    'n03160309', 'n03179701', 'n03180011', 'n03187595', 'n03188531',
+    'n03196217', 'n03197337', 'n03201208', 'n03207743', 'n03207941',
+    'n03208938', 'n03216828', 'n03218198', 'n03220513', 'n03223299',
+    'n03240683', 'n03249569', 'n03250847', 'n03255030', 'n03259280',
+    'n03271574', 'n03272010', 'n03272562', 'n03290653', 'n03291819',
+    'n03297495', 'n03314780', 'n03325584', 'n03337140', 'n03344393',
+    'n03345487', 'n03347037', 'n03355925', 'n03372029', 'n03376595',
+    'n03379051', 'n03384352', 'n03388043', 'n03388183', 'n03388549',
+    'n03393912', 'n03394916', 'n03400231', 'n03404251', 'n03417042',
+    'n03424325', 'n03425413', 'n03443371', 'n03444034', 'n03445777',
+    'n03445924', 'n03447447', 'n03447721', 'n03450230', 'n03452741',
+    'n03457902', 'n03459775', 'n03461385', 'n03467068', 'n03476684',
+    'n03476991', 'n03478589', 'n03481172', 'n03482405', 'n03483316',
+    'n03485407', 'n03485794', 'n03492542', 'n03494278', 'n03495258',
+    'n03496892', 'n03498962', 'n03527444', 'n03529860', 'n03530642',
+    'n03532672', 'n03534580', 'n03535780', 'n03538406', 'n03544143',
+    'n03584254', 'n03584829', 'n03590841', 'n03594734', 'n03594945',
+    'n03595614', 'n03598930', 'n03599486', 'n03602883', 'n03617480',
+    'n03623198', 'n03627232', 'n03630383', 'n03633091', 'n03637318',
+    'n03642806', 'n03649909', 'n03657121', 'n03658185', 'n03661043',
+    'n03662601', 'n03666591', 'n03670208', 'n03673027', 'n03676483',
+    'n03680355', 'n03690938', 'n03691459', 'n03692522', 'n03697007',
+    'n03706229', 'n03709823', 'n03710193', 'n03710637', 'n03710721',
+    'n03717622', 'n03720891', 'n03721384', 'n03724870', 'n03729826',
+    'n03733131', 'n03733281', 'n03733805', 'n03742115', 'n03743016',
+    'n03759954', 'n03761084', 'n03763968', 'n03764736', 'n03769881',
+    'n03770439', 'n03770679', 'n03773504', 'n03775071', 'n03775546',
+    'n03776460', 'n03777568', 'n03777754', 'n03781244', 'n03782006',
+    'n03785016', 'n03786901', 'n03787032', 'n03788195', 'n03788365',
+    'n03791053', 'n03792782', 'n03792972', 'n03793489', 'n03794056',
+    'n03796401', 'n03803284', 'n03804744', 'n03814639', 'n03814906',
+    'n03825788', 'n03832673', 'n03837869', 'n03838899', 'n03840681',
+    'n03841143', 'n03843555', 'n03854065', 'n03857828', 'n03866082',
+    'n03868242', 'n03868863', 'n03871628', 'n03873416', 'n03874293',
+    'n03874599', 'n03876231', 'n03877472', 'n03877845', 'n03884397',
+    'n03887697', 'n03888257', 'n03888605', 'n03891251', 'n03891332',
+    'n03895866', 'n03899768', 'n03902125', 'n03903868', 'n03908618',
+    'n03908714', 'n03916031', 'n03920288', 'n03924679', 'n03929660',
+    'n03929855', 'n03930313', 'n03930630', 'n03933933', 'n03935335',
+    'n03937543', 'n03938244', 'n03942813', 'n03944341', 'n03947888',
+    'n03950228', 'n03954731', 'n03956157', 'n03958227', 'n03961711',
+    'n03967562', 'n03970156', 'n03976467', 'n03976657', 'n03977966',
+    'n03980874', 'n03982430', 'n03983396', 'n03991062', 'n03992509',
+    'n03995372', 'n03998194', 'n04004767', 'n04005630', 'n04008634',
+    'n04009552', 'n04019541', 'n04023962', 'n04026417', 'n04033901',
+    'n04033995', 'n04037443', 'n04039381', 'n04040759', 'n04041544',
+    'n04044716', 'n04049303', 'n04065272', 'n04067472', 'n04069434',
+    'n04070727', 'n04074963', 'n04081281', 'n04086273', 'n04090263',
+    'n04099969', 'n04111531', 'n04116512', 'n04118538', 'n04118776',
+    'n04120489', 'n04125021', 'n04127249', 'n04131690', 'n04133789',
+    'n04136333', 'n04141076', 'n04141327', 'n04141975', 'n04146614',
+    'n04147183', 'n04149813', 'n04152593', 'n04153751', 'n04154565',
+    'n04162706', 'n04179913', 'n04192698', 'n04200800', 'n04201297',
+    'n04204238', 'n04204347', 'n04208210', 'n04209133', 'n04209239',
+    'n04228054', 'n04229816', 'n04235860', 'n04238763', 'n04239074',
+    'n04243546', 'n04251144', 'n04252077', 'n04252225', 'n04254120',
+    'n04254680', 'n04254777', 'n04258138', 'n04259630', 'n04263257',
+    'n04264628', 'n04265275', 'n04266014', 'n04270147', 'n04273569',
+    'n04275548', 'n04277352', 'n04285008', 'n04286575', 'n04296562',
+    'n04310018', 'n04311004', 'n04311174', 'n04317175', 'n04325704',
+    'n04326547', 'n04328186', 'n04330267', 'n04332243', 'n04335435',
+    'n04336792', 'n04344873', 'n04346328', 'n04347754', 'n04350905',
+    'n04355338', 'n04355933', 'n04356056', 'n04357314', 'n04366367',
+    'n04367480', 'n04370456', 'n04371430', 'n04371774', 'n04372370',
+    'n04376876', 'n04380533', 'n04389033', 'n04392985', 'n04398044',
+    'n04399382', 'n04404412', 'n04409515', 'n04417672', 'n04418357',
+    'n04423845', 'n04428191', 'n04429376', 'n04435653', 'n04442312',
+    'n04443257', 'n04447861', 'n04456115', 'n04458633', 'n04461696',
+    'n04462240', 'n04465501', 'n04467665', 'n04476259', 'n04479046',
+    'n04482393', 'n04483307', 'n04485082', 'n04486054', 'n04487081',
+    'n04487394', 'n04493381', 'n04501370', 'n04505470', 'n04507155',
+    'n04509417', 'n04515003', 'n04517823', 'n04522168', 'n04523525',
+    'n04525038', 'n04525305', 'n04532106', 'n04532670', 'n04536866',
+    'n04540053', 'n04542943', 'n04548280', 'n04548362', 'n04550184',
+    'n04552348', 'n04553703', 'n04554684', 'n04557648', 'n04560804',
+    'n04562935', 'n04579145', 'n04579432', 'n04584207', 'n04589890',
+    'n04590129', 'n04591157', 'n04591713', 'n04592741', 'n04596742',
+    'n04597913', 'n04599235', 'n04604644', 'n04606251', 'n04612504',
+    'n04613696', 'n06359193', 'n06596364', 'n06785654', 'n06794110',
+    'n06874185', 'n07248320', 'n07565083', 'n07579787', 'n07583066',
+    'n07584110', 'n07590611', 'n07613480', 'n07614500', 'n07615774',
+    'n07684084', 'n07693725', 'n07695742', 'n07697313', 'n07697537',
+    'n07711569', 'n07714571', 'n07714990', 'n07715103', 'n07716358',
+    'n07716906', 'n07717410', 'n07717556', 'n07718472', 'n07718747',
+    'n07720875', 'n07730033', 'n07734744', 'n07742313', 'n07745940',
+    'n07747607', 'n07749582', 'n07753113', 'n07753275', 'n07753592',
+    'n07754684', 'n07760859', 'n07768694', 'n07802026', 'n07831146',
+    'n07836838', 'n07860988', 'n07871810', 'n07873807', 'n07875152',
+    'n07880968', 'n07892512', 'n07920052', 'n07930864', 'n07932039',
+    'n09193705', 'n09229709', 'n09246464', 'n09256479', 'n09288635',
+    'n09332890', 'n09399592', 'n09421951', 'n09428293', 'n09468604',
+    'n09472597', 'n09835506', 'n10148035', 'n10565667', 'n11879895',
+    'n11939491', 'n12057211', 'n12144580', 'n12267677', 'n12620546',
+    'n12768682', 'n12985857', 'n12998815', 'n13037406', 'n13040303',
+    'n13044778', 'n13052670', 'n13054560', 'n13133613', 'n15075141'
+]
+
+imagenet_a_wnids = [
+    'n01498041', 'n01531178', 'n01534433', 'n01558993', 'n01580077',
+    'n01614925', 'n01616318', 'n01631663', 'n01641577', 'n01669191',
+    'n01677366', 'n01687978', 'n01694178', 'n01698640', 'n01735189',
+    'n01770081', 'n01770393', 'n01774750', 'n01784675', 'n01819313',
+    'n01820546', 'n01833805', 'n01843383', 'n01847000', 'n01855672',
+    'n01882714', 'n01910747', 'n01914609', 'n01924916', 'n01944390',
+    'n01985128', 'n01986214', 'n02007558', 'n02009912', 'n02037110',
+    'n02051845', 'n02077923', 'n02085620', 'n02099601', 'n02106550',
+    'n02106662', 'n02110958', 'n02119022', 'n02123394', 'n02127052',
+    'n02129165', 'n02133161', 'n02137549', 'n02165456', 'n02174001',
+    'n02177972', 'n02190166', 'n02206856', 'n02219486', 'n02226429',
+    'n02231487', 'n02233338', 'n02236044', 'n02259212', 'n02268443',
+    'n02279972', 'n02280649', 'n02281787', 'n02317335', 'n02325366',
+    'n02346627', 'n02356798', 'n02361337', 'n02410509', 'n02445715',
+    'n02454379', 'n02486410', 'n02492035', 'n02504458', 'n02655020',
+    'n02669723', 'n02672831', 'n02676566', 'n02690373', 'n02701002',
+    'n02730930', 'n02777292', 'n02782093', 'n02787622', 'n02793495',
+    'n02797295', 'n02802426', 'n02814860', 'n02815834', 'n02837789',
+    'n02879718', 'n02883205', 'n02895154', 'n02906734', 'n02948072',
+    'n02951358', 'n02980441', 'n02992211', 'n02999410', 'n03014705',
+    'n03026506', 'n03124043', 'n03125729', 'n03187595', 'n03196217',
+    'n03223299', 'n03250847', 'n03255030', 'n03291819', 'n03325584',
+    'n03355925', 'n03384352', 'n03388043', 'n03417042', 'n03443371',
+    'n03444034', 'n03445924', 'n03452741', 'n03483316', 'n03584829',
+    'n03590841', 'n03594945', 'n03617480', 'n03666591', 'n03670208',
+    'n03717622', 'n03720891', 'n03721384', 'n03724870', 'n03775071',
+    'n03788195', 'n03804744', 'n03837869', 'n03840681', 'n03854065',
+    'n03888257', 'n03891332', 'n03935335', 'n03982430', 'n04019541',
+    'n04033901', 'n04039381', 'n04067472', 'n04086273', 'n04099969',
+    'n04118538', 'n04131690', 'n04133789', 'n04141076', 'n04146614',
+    'n04147183', 'n04179913', 'n04208210', 'n04235860', 'n04252077',
+    'n04252225', 'n04254120', 'n04270147', 'n04275548', 'n04310018',
+    'n04317175', 'n04344873', 'n04347754', 'n04355338', 'n04366367',
+    'n04376876', 'n04389033', 'n04399382', 'n04442312', 'n04456115',
+    'n04482393', 'n04507155', 'n04509417', 'n04532670', 'n04540053',
+    'n04554684', 'n04562935', 'n04591713', 'n04606251', 'n07583066',
+    'n07695742', 'n07697313', 'n07697537', 'n07714990', 'n07718472',
+    'n07720875', 'n07734744', 'n07749582', 'n07753592', 'n07760859',
+    'n07768694', 'n07831146', 'n09229709', 'n09246464', 'n09472597',
+    'n09835506', 'n11879895', 'n12057211', 'n12144580', 'n12267677'
+]
+
+imagenet_a_mask = [wnid in set(imagenet_a_wnids) for wnid in all_wnids]
+
+imagenet_r_wnids = {
+    'n01443537', 'n01484850', 'n01494475', 'n01498041', 'n01514859',
+    'n01518878', 'n01531178', 'n01534433', 'n01614925', 'n01616318',
+    'n01630670', 'n01632777', 'n01644373', 'n01677366', 'n01694178',
+    'n01748264', 'n01770393', 'n01774750', 'n01784675', 'n01806143',
+    'n01820546', 'n01833805', 'n01843383', 'n01847000', 'n01855672',
+    'n01860187', 'n01882714', 'n01910747', 'n01944390', 'n01983481',
+    'n01986214', 'n02007558', 'n02009912', 'n02051845', 'n02056570',
+    'n02066245', 'n02071294', 'n02077923', 'n02085620', 'n02086240',
+    'n02088094', 'n02088238', 'n02088364', 'n02088466', 'n02091032',
+    'n02091134', 'n02092339', 'n02094433', 'n02096585', 'n02097298',
+    'n02098286', 'n02099601', 'n02099712', 'n02102318', 'n02106030',
+    'n02106166', 'n02106550', 'n02106662', 'n02108089', 'n02108915',
+    'n02109525', 'n02110185', 'n02110341', 'n02110958', 'n02112018',
+    'n02112137', 'n02113023', 'n02113624', 'n02113799', 'n02114367',
+    'n02117135', 'n02119022', 'n02123045', 'n02128385', 'n02128757',
+    'n02129165', 'n02129604', 'n02130308', 'n02134084', 'n02138441',
+    'n02165456', 'n02190166', 'n02206856', 'n02219486', 'n02226429',
+    'n02233338', 'n02236044', 'n02268443', 'n02279972', 'n02317335',
+    'n02325366', 'n02346627', 'n02356798', 'n02363005', 'n02364673',
+    'n02391049', 'n02395406', 'n02398521', 'n02410509', 'n02423022',
+    'n02437616', 'n02445715', 'n02447366', 'n02480495', 'n02480855',
+    'n02481823', 'n02483362', 'n02486410', 'n02510455', 'n02526121',
+    'n02607072', 'n02655020', 'n02672831', 'n02701002', 'n02749479',
+    'n02769748', 'n02793495', 'n02797295', 'n02802426', 'n02808440',
+    'n02814860', 'n02823750', 'n02841315', 'n02843684', 'n02883205',
+    'n02906734', 'n02909870', 'n02939185', 'n02948072', 'n02950826',
+    'n02951358', 'n02966193', 'n02980441', 'n02992529', 'n03124170',
+    'n03272010', 'n03345487', 'n03372029', 'n03424325', 'n03452741',
+    'n03467068', 'n03481172', 'n03494278', 'n03495258', 'n03498962',
+    'n03594945', 'n03602883', 'n03630383', 'n03649909', 'n03676483',
+    'n03710193', 'n03773504', 'n03775071', 'n03888257', 'n03930630',
+    'n03947888', 'n04086273', 'n04118538', 'n04133789', 'n04141076',
+    'n04146614', 'n04147183', 'n04192698', 'n04254680', 'n04266014',
+    'n04275548', 'n04310018', 'n04325704', 'n04347754', 'n04389033',
+    'n04409515', 'n04465501', 'n04487394', 'n04522168', 'n04536866',
+    'n04552348', 'n04591713', 'n07614500', 'n07693725', 'n07695742',
+    'n07697313', 'n07697537', 'n07714571', 'n07714990', 'n07718472',
+    'n07720875', 'n07734744', 'n07742313', 'n07745940', 'n07749582',
+    'n07753275', 'n07753592', 'n07768694', 'n07873807', 'n07880968',
+    'n07920052', 'n09472597', 'n09835506', 'n10565667', 'n12267677'
+}
+
+imagenet_r_mask = [wnid in imagenet_r_wnids for wnid in all_wnids]

classification/dataset/imagenet_real.py

+# --------------------------------------------------------
+# EVA: Exploring the Limits of Masked Visual Representation Learning at Scale (https://arxiv.org/abs/2211.07636)
+# Github source: https://github.com/baaivision/EVA
+# Copyright (c) 2022 Beijing Academy of Artificial Intelligence (BAAI)
+# Licensed under The MIT License [see LICENSE for details]
+# By Yuxin Fang
+# Based on timm, DINO, DeiT and BEiT codebases
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# https://github.com/facebookresearch/deit
+# https://github.com/facebookresearch/dino
+# https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------'
+
+import json
+import os
+
+import numpy as np
+
+
+class RealLabelsImagenet:
+
+    def __init__(self, filenames, real_json='real.json', topk=(1, 5)):
+        with open(real_json) as real_labels:
+            real_labels = json.load(real_labels)
+            real_labels = {f'ILSVRC2012_val_{i + 1:08d}.JPEG': labels for i, labels in enumerate(real_labels)}
+        self.real_labels = real_labels
+        self.filenames = filenames
+        assert len(self.filenames) == len(self.real_labels)
+        self.topk = topk
+        self.is_correct = {k: [] for k in topk}
+        self.sample_idx = 0
+
+    def add_result(self, output):
+        maxk = max(self.topk)
+        _, pred_batch = output.topk(maxk, 1, True, True)
+        pred_batch = pred_batch.cpu().numpy()
+        for pred in pred_batch:
+            filename = self.filenames[self.sample_idx]
+            filename = os.path.basename(filename)
+            if self.real_labels[filename]:
+                for k in self.topk:
+                    self.is_correct[k].append(
+                        any([p in self.real_labels[filename] for p in pred[:k]]))
+            self.sample_idx += 1
+
+    def get_accuracy(self, k=None):
+        if k is None:
+            return {k: float(np.mean(self.is_correct[k] for k in self.topk))}
+        else:
+            return float(np.mean(self.is_correct[k])) * 100

classification/dataset/imagenetv2.py

+"""Code from https://github.com/mlfoundations/wise-ft/blob/master/src/datasets/imagenetv2.py
+Thanks to the authors of wise-ft."""
+import pathlib
+import shutil
+import tarfile
+
+import requests
+from PIL import Image
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+URLS = {'matched-frequency': 'https://imagenetv2public.s3-us-west-2.amazonaws.com/imagenetv2-matched-frequency.tar.gz',
+        'threshold-0.7': 'https://imagenetv2public.s3-us-west-2.amazonaws.com/imagenetv2-threshold0.7.tar.gz',
+        'top-images': 'https://imagenetv2public.s3-us-west-2.amazonaws.com/imagenetv2-top-images.tar.gz',
+        'val': 'https://imagenetv2public.s3-us-west-2.amazonaws.com/imagenet_validation.tar.gz'}
+
+FNAMES = {'matched-frequency': 'imagenetv2-matched-frequency-format-val',
+          'threshold-0.7': 'imagenetv2-threshold0.7-format-val',
+          'top-images': 'imagenetv2-top-images-format-val',
+          'val': 'imagenet_validation'}
+
+V2_DATASET_SIZE = 10000
+VAL_DATASET_SIZE = 50000
+
+
+class ImageNetV2Dataset(Dataset):
+    def __init__(self, variant='matched-frequency', transform=None, location='.'):
+        self.dataset_root = pathlib.Path(f'{location}/ImageNetV2-{variant}/')
+        self.tar_root = pathlib.Path(f'{location}/ImageNetV2-{variant}.tar.gz')
+        self.fnames = list(self.dataset_root.glob('**/*.jpeg'))
+        self.transform = transform
+        assert variant in URLS, f'unknown V2 Variant: {variant}'
+        if not self.dataset_root.exists() or len(self.fnames) != V2_DATASET_SIZE:
+            if not self.tar_root.exists():
+                print(f'Dataset {variant} not found on disk, downloading....')
+                response = requests.get(URLS[variant], stream=True)
+                total_size_in_bytes = int(response.headers.get('content-length', 0))
+                block_size = 1024  # 1 Kibibyte
+                progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
+                with open(self.tar_root, 'wb') as f:
+                    for data in response.iter_content(block_size):
+                        progress_bar.update(len(data))
+                        f.write(data)
+                progress_bar.close()
+                if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
+                    assert False, f'Downloading from {URLS[variant]} failed'
+            print('Extracting....')
+            tarfile.open(self.tar_root).extractall(f'{location}')
+            shutil.move(f'{location}/{FNAMES[variant]}', self.dataset_root)
+            self.fnames = list(self.dataset_root.glob('**/*.jpeg'))
+
+    def __len__(self):
+        return len(self.fnames)
+
+    def __getitem__(self, i):
+        img, label = Image.open(self.fnames[i]), int(self.fnames[i].parent.name)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label

classification/dataset/samplers.py

+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import math
+import os
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.utils.data.sampler import Sampler
+
+
+class SubsetRandomSampler(torch.utils.data.Sampler):
+    """Samples elements randomly from a given list of indices, without
+    replacement.
+
+    Arguments:
+        indices (sequence): a sequence of indices
+    """
+
+    def __init__(self, indices):
+        self.epoch = 0
+        self.indices = indices
+
+    def __iter__(self):
+        return (self.indices[i] for i in torch.randperm(len(self.indices)))
+
+    def __len__(self):
+        return len(self.indices)
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+
+class NodeDistributedSampler(Sampler):
+    """Sampler that restricts data loading to a subset of the dataset.
+    It is especially useful in conjunction with
+    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
+    process can pass a DistributedSampler instance as a DataLoader sampler,
+    and load a subset of the original dataset that is exclusive to it.
+    .. note::
+        Dataset is assumed to be of constant size.
+    Arguments:
+        dataset: Dataset used for sampling.
+        num_replicas (optional): Number of processes participating in
+            distributed training.
+        rank (optional): Rank of the current process within num_replicas.
+    """
+
+    def __init__(self,
+                 dataset,
+                 num_replicas=None,
+                 rank=None,
+                 local_rank=None,
+                 local_size=None):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError(
+                    'Requires distributed package to be available')
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError(
+                    'Requires distributed package to be available')
+            rank = dist.get_rank()
+        if local_rank is None:
+            local_rank = int(os.environ.get('LOCAL_RANK', 0))
+        if local_size is None:
+            local_size = int(os.environ.get('LOCAL_SIZE', 1))
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.num_parts = local_size
+        self.rank = rank
+        self.local_rank = local_rank
+        self.epoch = 0
+        self.num_samples = int(
+            math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+
+        self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+
+        t = torch.Generator()
+        t.manual_seed(0)
+
+        indices = torch.randperm(len(self.dataset), generator=t).tolist()
+        # indices = range(len(self.dataset))
+        indices = [i for i in indices if i % self.num_parts == self.local_rank]
+
+        # add extra samples to make it evenly divisible
+        indices += indices[:(self.total_size_parts - len(indices))]
+        assert len(indices) == self.total_size_parts
+
+        # subsample
+        indices = indices[self.rank // self.num_parts:self.
+                          total_size_parts:self.num_replicas // self.num_parts]
+
+        index = torch.randperm(len(indices), generator=g).tolist()
+        indices = list(np.array(indices)[index])
+
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

classification/logger.py

+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2022 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import functools
+import logging
+import os
+import sys
+
+from termcolor import colored
+
+
+@functools.lru_cache()
+def create_logger(output_dir, dist_rank=0, name=''):
+    # create logger
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+    logger.propagate = False
+
+    # create formatter
+    fmt = '[%(asctime)s %(name)s] (%(filename)s %(lineno)d): %(levelname)s %(message)s'
+    color_fmt = colored('[%(asctime)s %(name)s]', 'green') + \
+        colored('(%(filename)s %(lineno)d)', 'yellow') + \
+        ': %(levelname)s %(message)s'
+
+    # create console handlers for master process
+    if dist_rank == 0:
+        console_handler = logging.StreamHandler(sys.stdout)
+        console_handler.setLevel(logging.DEBUG)
+        console_handler.setFormatter(
+            logging.Formatter(fmt=color_fmt, datefmt='%Y-%m-%d %H:%M:%S'))
+        logger.addHandler(console_handler)
+
+    # create file handlers
+    file_handler = logging.FileHandler(os.path.join(
+        output_dir, f'log_rank{dist_rank}.txt'),
+                                       mode='a')
+    file_handler.setLevel(logging.DEBUG)
+    file_handler.setFormatter(
+        logging.Formatter(fmt=fmt, datefmt='%Y-%m-%d %H:%M:%S'))
+    logger.addHandler(file_handler)
+
+    return logger