From 055c506510dd456f4a5f6cc52655a35ebbb1f87b Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Wed, 19 Oct 2022 18:08:24 +0000 Subject: [PATCH 01/47] isort & black --- config.py | 89 ++-- data/__init__.py | 2 +- data/build.py | 95 ++-- data/cached_image_folder.py | 97 +++- data/data_simmim_ft.py | 64 ++- data/data_simmim_pt.py | 83 +-- data/imagenet22k_dataset.py | 12 +- data/zipreader.py | 49 +- kernels/window_process/setup.py | 18 +- kernels/window_process/unit_test.py | 183 +++++-- kernels/window_process/window_process.py | 38 +- logger.py | 25 +- lr_scheduler.py | 88 +-- main.py | 323 +++++++---- main_moe.py | 360 +++++++++---- main_simmim_ft.py | 258 ++++++--- main_simmim_pt.py | 170 ++++-- models/__init__.py | 2 +- models/build.py | 184 ++++--- models/simmim.py | 107 ++-- models/swin_mlp.py | 296 ++++++++--- models/swin_transformer.py | 405 ++++++++++---- models/swin_transformer_moe.py | 649 ++++++++++++++++------- models/swin_transformer_v2.py | 418 +++++++++++---- optimizer.py | 123 +++-- utils.py | 184 ++++--- utils_moe.py | 240 ++++++--- utils_simmim.py | 139 +++-- 28 files changed, 3210 insertions(+), 1491 deletions(-) diff --git a/config.py b/config.py index 1671ec34..e5ba8761 100644 --- a/config.py +++ b/config.py @@ -6,13 +6,14 @@ # --------------------------------------------------------' import os + import yaml from yacs.config import CfgNode as CN _C = CN() # Base config files -_C.BASE = [''] +_C.BASE = [""] # ----------------------------------------------------------------------------- # Data settings @@ -21,18 +22,18 @@ # Batch size for a single GPU, could be overwritten by command line argument _C.DATA.BATCH_SIZE = 128 # Path to dataset, could be overwritten by command line argument -_C.DATA.DATA_PATH = '' +_C.DATA.DATA_PATH = "" # Dataset name -_C.DATA.DATASET = 'imagenet' +_C.DATA.DATASET = "imagenet" # Input image size _C.DATA.IMG_SIZE = 224 # Interpolation to resize image (random, bilinear, bicubic) -_C.DATA.INTERPOLATION = 'bicubic' +_C.DATA.INTERPOLATION = "bicubic" # Use zipped dataset instead of folder dataset # could be overwritten by command line argument _C.DATA.ZIP_MODE = False # Cache Data in Memory, could be overwritten by command line argument -_C.DATA.CACHE_MODE = 'part' +_C.DATA.CACHE_MODE = "part" # Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU. _C.DATA.PIN_MEMORY = True # Number of data loading threads @@ -48,14 +49,14 @@ # ----------------------------------------------------------------------------- _C.MODEL = CN() # Model type -_C.MODEL.TYPE = 'swin' +_C.MODEL.TYPE = "swin" # Model name -_C.MODEL.NAME = 'swin_tiny_patch4_window7_224' +_C.MODEL.NAME = "swin_tiny_patch4_window7_224" # Pretrained weight from checkpoint, could be imagenet22k pretrained weight # could be overwritten by command line argument -_C.MODEL.PRETRAINED = '' +_C.MODEL.PRETRAINED = "" # Checkpoint to resume, could be overwritten by command line argument -_C.MODEL.RESUME = '' +_C.MODEL.RESUME = "" # Number of classes, overwritten in data preparation _C.MODEL.NUM_CLASSES = 1000 # Dropout rate @@ -73,7 +74,7 @@ _C.MODEL.SWIN.DEPTHS = [2, 2, 6, 2] _C.MODEL.SWIN.NUM_HEADS = [3, 6, 12, 24] _C.MODEL.SWIN.WINDOW_SIZE = 7 -_C.MODEL.SWIN.MLP_RATIO = 4. +_C.MODEL.SWIN.MLP_RATIO = 4.0 _C.MODEL.SWIN.QKV_BIAS = True _C.MODEL.SWIN.QK_SCALE = None _C.MODEL.SWIN.APE = False @@ -87,7 +88,7 @@ _C.MODEL.SWINV2.DEPTHS = [2, 2, 6, 2] _C.MODEL.SWINV2.NUM_HEADS = [3, 6, 12, 24] _C.MODEL.SWINV2.WINDOW_SIZE = 7 -_C.MODEL.SWINV2.MLP_RATIO = 4. +_C.MODEL.SWINV2.MLP_RATIO = 4.0 _C.MODEL.SWINV2.QKV_BIAS = True _C.MODEL.SWINV2.APE = False _C.MODEL.SWINV2.PATCH_NORM = True @@ -101,7 +102,7 @@ _C.MODEL.SWIN_MOE.DEPTHS = [2, 2, 6, 2] _C.MODEL.SWIN_MOE.NUM_HEADS = [3, 6, 12, 24] _C.MODEL.SWIN_MOE.WINDOW_SIZE = 7 -_C.MODEL.SWIN_MOE.MLP_RATIO = 4. +_C.MODEL.SWIN_MOE.MLP_RATIO = 4.0 _C.MODEL.SWIN_MOE.QKV_BIAS = True _C.MODEL.SWIN_MOE.QK_SCALE = None _C.MODEL.SWIN_MOE.APE = False @@ -131,7 +132,7 @@ _C.MODEL.SWIN_MLP.DEPTHS = [2, 2, 6, 2] _C.MODEL.SWIN_MLP.NUM_HEADS = [3, 6, 12, 24] _C.MODEL.SWIN_MLP.WINDOW_SIZE = 7 -_C.MODEL.SWIN_MLP.MLP_RATIO = 4. +_C.MODEL.SWIN_MLP.MLP_RATIO = 4.0 _C.MODEL.SWIN_MLP.APE = False _C.MODEL.SWIN_MLP.PATCH_NORM = True @@ -165,7 +166,7 @@ # LR scheduler _C.TRAIN.LR_SCHEDULER = CN() -_C.TRAIN.LR_SCHEDULER.NAME = 'cosine' +_C.TRAIN.LR_SCHEDULER.NAME = "cosine" # Epoch interval to decay LR, used in StepLRScheduler _C.TRAIN.LR_SCHEDULER.DECAY_EPOCHS = 30 # LR decay rate, used in StepLRScheduler @@ -178,7 +179,7 @@ # Optimizer _C.TRAIN.OPTIMIZER = CN() -_C.TRAIN.OPTIMIZER.NAME = 'adamw' +_C.TRAIN.OPTIMIZER.NAME = "adamw" # Optimizer Epsilon _C.TRAIN.OPTIMIZER.EPS = 1e-8 # Optimizer Betas @@ -200,11 +201,11 @@ # Color jitter factor _C.AUG.COLOR_JITTER = 0.4 # Use AutoAugment policy. "v0" or "original" -_C.AUG.AUTO_AUGMENT = 'rand-m9-mstd0.5-inc1' +_C.AUG.AUTO_AUGMENT = "rand-m9-mstd0.5-inc1" # Random erase prob _C.AUG.REPROB = 0.25 # Random erase mode -_C.AUG.REMODE = 'pixel' +_C.AUG.REMODE = "pixel" # Random erase count _C.AUG.RECOUNT = 1 # Mixup alpha, mixup enabled if > 0 @@ -218,7 +219,7 @@ # Probability of switching to cutmix when both mixup and cutmix enabled _C.AUG.MIXUP_SWITCH_PROB = 0.5 # How to apply mixup/cutmix params. Per "batch", "pair", or "elem" -_C.AUG.MIXUP_MODE = 'batch' +_C.AUG.MIXUP_MODE = "batch" # ----------------------------------------------------------------------------- # Testing settings @@ -239,11 +240,11 @@ # Enable Pytorch automatic mixed precision (amp). _C.AMP_ENABLE = True # [Deprecated] Mixed precision opt level of apex, if O0, no apex amp is used ('O0', 'O1', 'O2') -_C.AMP_OPT_LEVEL = '' +_C.AMP_OPT_LEVEL = "" # Path to output folder, overwritten by command line argument -_C.OUTPUT = '' +_C.OUTPUT = "" # Tag of experiment, overwritten by command line argument -_C.TAG = 'default' +_C.TAG = "default" # Frequency to save checkpoint _C.SAVE_FREQ = 1 # Frequency to logging info @@ -263,15 +264,15 @@ def _update_config_from_file(config, cfg_file): config.defrost() - with open(cfg_file, 'r') as f: + with open(cfg_file, "r") as f: yaml_cfg = yaml.load(f, Loader=yaml.FullLoader) - for cfg in yaml_cfg.setdefault('BASE', ['']): + for cfg in yaml_cfg.setdefault("BASE", [""]): if cfg: _update_config_from_file( config, os.path.join(os.path.dirname(cfg_file), cfg) ) - print('=> merge config from {}'.format(cfg_file)) + print("=> merge config from {}".format(cfg_file)) config.merge_from_file(cfg_file) config.freeze() @@ -284,53 +285,53 @@ def update_config(config, args): config.merge_from_list(args.opts) def _check_args(name): - if hasattr(args, name) and eval(f'args.{name}'): + if hasattr(args, name) and eval(f"args.{name}"): return True return False # merge from specific arguments - if _check_args('batch_size'): + if _check_args("batch_size"): config.DATA.BATCH_SIZE = args.batch_size - if _check_args('data_path'): + if _check_args("data_path"): config.DATA.DATA_PATH = args.data_path - if _check_args('zip'): + if _check_args("zip"): config.DATA.ZIP_MODE = True - if _check_args('cache_mode'): + if _check_args("cache_mode"): config.DATA.CACHE_MODE = args.cache_mode - if _check_args('pretrained'): + if _check_args("pretrained"): config.MODEL.PRETRAINED = args.pretrained - if _check_args('resume'): + if _check_args("resume"): config.MODEL.RESUME = args.resume - if _check_args('accumulation_steps'): + if _check_args("accumulation_steps"): config.TRAIN.ACCUMULATION_STEPS = args.accumulation_steps - if _check_args('use_checkpoint'): + if _check_args("use_checkpoint"): config.TRAIN.USE_CHECKPOINT = True - if _check_args('amp_opt_level'): + if _check_args("amp_opt_level"): print("[warning] Apex amp has been deprecated, please use pytorch amp instead!") - if args.amp_opt_level == 'O0': + if args.amp_opt_level == "O0": config.AMP_ENABLE = False - if _check_args('disable_amp'): + if _check_args("disable_amp"): config.AMP_ENABLE = False - if _check_args('output'): + if _check_args("output"): config.OUTPUT = args.output - if _check_args('tag'): + if _check_args("tag"): config.TAG = args.tag - if _check_args('eval'): + if _check_args("eval"): config.EVAL_MODE = True - if _check_args('throughput'): + if _check_args("throughput"): config.THROUGHPUT_MODE = True # [SimMIM] - if _check_args('enable_amp'): + if _check_args("enable_amp"): config.ENABLE_AMP = args.enable_amp # for acceleration - if _check_args('fused_window_process'): + if _check_args("fused_window_process"): config.FUSED_WINDOW_PROCESS = True - if _check_args('fused_layernorm'): + if _check_args("fused_layernorm"): config.FUSED_LAYERNORM = True ## Overwrite optimizer if not None, currently we use it for [fused_adam, fused_lamb] - if _check_args('optim'): + if _check_args("optim"): config.TRAIN.OPTIMIZER.NAME = args.optim # set local rank for distributed training diff --git a/data/__init__.py b/data/__init__.py index 5baad7ed..a8cfa99a 100644 --- a/data/__init__.py +++ b/data/__init__.py @@ -1,6 +1,6 @@ from .build import build_loader as _build_loader -from .data_simmim_pt import build_loader_simmim from .data_simmim_ft import build_loader_finetune +from .data_simmim_pt import build_loader_simmim def build_loader(config, simmim=False, is_pretrain=False): diff --git a/data/build.py b/data/build.py index 5799f253..85a93e98 100644 --- a/data/build.py +++ b/data/build.py @@ -6,13 +6,13 @@ # -------------------------------------------------------- import os -import torch + import numpy as np +import torch import torch.distributed as dist -from torchvision import datasets, transforms +from timm.data import Mixup, create_transform from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD -from timm.data import Mixup -from timm.data import create_transform +from torchvision import datasets, transforms from .cached_image_folder import CachedImageFolder from .imagenet22k_dataset import IN22KDATASET @@ -21,19 +21,17 @@ try: from torchvision.transforms import InterpolationMode - def _pil_interp(method): - if method == 'bicubic': + if method == "bicubic": return InterpolationMode.BICUBIC - elif method == 'lanczos': + elif method == "lanczos": return InterpolationMode.LANCZOS - elif method == 'hamming': + elif method == "hamming": return InterpolationMode.HAMMING else: # default bilinear, do we want to allow nearest? return InterpolationMode.BILINEAR - import timm.data.transforms as timm_transforms timm_transforms._pil_interp = _pil_interp @@ -43,15 +41,21 @@ def _pil_interp(method): def build_loader(config): config.defrost() - dataset_train, config.MODEL.NUM_CLASSES = build_dataset(is_train=True, config=config) + dataset_train, config.MODEL.NUM_CLASSES = build_dataset( + is_train=True, config=config + ) config.freeze() - print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build train dataset") + print( + f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build train dataset" + ) dataset_val, _ = build_dataset(is_train=False, config=config) - print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build val dataset") + print( + f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build val dataset" + ) num_tasks = dist.get_world_size() global_rank = dist.get_rank() - if config.DATA.ZIP_MODE and config.DATA.CACHE_MODE == 'part': + 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: @@ -67,7 +71,8 @@ def build_loader(config): ) data_loader_train = torch.utils.data.DataLoader( - dataset_train, sampler=sampler_train, + dataset_train, + sampler=sampler_train, batch_size=config.DATA.BATCH_SIZE, num_workers=config.DATA.NUM_WORKERS, pin_memory=config.DATA.PIN_MEMORY, @@ -75,41 +80,57 @@ def build_loader(config): ) data_loader_val = torch.utils.data.DataLoader( - dataset_val, sampler=sampler_val, + 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 + drop_last=False, ) # 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 + mixup_active = ( + config.AUG.MIXUP > 0 + or config.AUG.CUTMIX > 0.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) + 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, data_loader_train, data_loader_val, mixup_fn def build_dataset(is_train, config): transform = build_transform(is_train, config) - if config.DATA.DATASET == 'imagenet': - prefix = 'train' if is_train else 'val' + if config.DATA.DATASET == "imagenet": + prefix = "train" if is_train else "val" if config.DATA.ZIP_MODE: ann_file = prefix + "_map.txt" prefix = prefix + ".zip@/" - dataset = CachedImageFolder(config.DATA.DATA_PATH, ann_file, prefix, transform, - cache_mode=config.DATA.CACHE_MODE if is_train else 'part') + dataset = CachedImageFolder( + config.DATA.DATA_PATH, + ann_file, + prefix, + transform, + cache_mode=config.DATA.CACHE_MODE if is_train else "part", + ) else: root = os.path.join(config.DATA.DATA_PATH, prefix) dataset = datasets.ImageFolder(root, transform=transform) nb_classes = 1000 - elif config.DATA.DATASET == 'imagenet22K': - prefix = 'ILSVRC2011fall_whole' + elif config.DATA.DATASET == "imagenet22K": + prefix = "ILSVRC2011fall_whole" if is_train: ann_file = prefix + "_map_train.txt" else: @@ -129,8 +150,12 @@ def build_transform(is_train, config): 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, + 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, @@ -139,7 +164,9 @@ def build_transform(is_train, config): if not resize_im: # replace RandomResizedCropAndInterpolation with # RandomCrop - transform.transforms[0] = transforms.RandomCrop(config.DATA.IMG_SIZE, padding=4) + transform.transforms[0] = transforms.RandomCrop( + config.DATA.IMG_SIZE, padding=4 + ) return transform t = [] @@ -147,14 +174,18 @@ def build_transform(is_train, config): if config.TEST.CROP: size = int((256 / 224) * config.DATA.IMG_SIZE) t.append( - transforms.Resize(size, interpolation=_pil_interp(config.DATA.INTERPOLATION)), + 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)) else: t.append( - transforms.Resize((config.DATA.IMG_SIZE, config.DATA.IMG_SIZE), - interpolation=_pil_interp(config.DATA.INTERPOLATION)) + transforms.Resize( + (config.DATA.IMG_SIZE, config.DATA.IMG_SIZE), + interpolation=_pil_interp(config.DATA.INTERPOLATION), + ) ) t.append(transforms.ToTensor()) diff --git a/data/cached_image_folder.py b/data/cached_image_folder.py index 7e1883b1..732d3236 100644 --- a/data/cached_image_folder.py +++ b/data/cached_image_folder.py @@ -8,11 +8,12 @@ import io import os import time + import torch.distributed as dist import torch.utils.data as data from PIL import Image -from .zipreader import is_zip_path, ZipReader +from .zipreader import ZipReader, is_zip_path def has_file_allowed_extension(filename, extensions): @@ -56,7 +57,7 @@ def make_dataset_with_ann(ann_file, img_prefix, extensions): 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')] + path_contents = [c for c in line_str.split("\t")] im_file_name = path_contents[0] class_index = int(path_contents[1]) @@ -89,21 +90,37 @@ class DatasetFolder(data.Dataset): 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"): + def __init__( + self, + root, + loader, + extensions, + ann_file="", + img_prefix="", + transform=None, + target_transform=None, + cache_mode="no", + ): # image folder mode - if ann_file == '': + 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) + 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))) + raise ( + RuntimeError( + "Found 0 files in subfolders of: " + + root + + "\n" + + "Supported extensions are: " + + ",".join(extensions) + ) + ) self.root = root self.loader = loader @@ -131,7 +148,9 @@ def init_cache(self): 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') + 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": @@ -162,17 +181,21 @@ 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))) + 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'] +IMG_EXTENSIONS = [".jpg", ".jpeg", ".png", ".ppm", ".bmp", ".pgm", ".tif"] def pil_loader(path): @@ -183,14 +206,15 @@ def pil_loader(path): data = ZipReader.read(path) img = Image.open(io.BytesIO(data)) else: - with open(path, 'rb') as f: + with open(path, "rb") as f: img = Image.open(f) - return img.convert('RGB') - return img.convert('RGB') + return img.convert("RGB") + return img.convert("RGB") def accimage_loader(path): import accimage + try: return accimage.Image(path) except IOError: @@ -200,7 +224,8 @@ def accimage_loader(path): def default_img_loader(path): from torchvision import get_image_backend - if get_image_backend() == 'accimage': + + if get_image_backend() == "accimage": return accimage_loader(path) else: return pil_loader(path) @@ -225,12 +250,26 @@ class CachedImageFolder(DatasetFolder): 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) + 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): diff --git a/data/data_simmim_ft.py b/data/data_simmim_ft.py index 1d44ae7d..ffd4137c 100644 --- a/data/data_simmim_ft.py +++ b/data/data_simmim_ft.py @@ -6,18 +6,20 @@ # -------------------------------------------------------- import os + import torch.distributed as dist -from torch.utils.data import DataLoader, DistributedSampler -from torchvision import datasets, transforms +from timm.data import Mixup, create_transform from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD -from timm.data import Mixup -from timm.data import create_transform from timm.data.transforms import _pil_interp +from torch.utils.data import DataLoader, DistributedSampler +from torchvision import datasets, transforms def build_loader_finetune(config): config.defrost() - dataset_train, config.MODEL.NUM_CLASSES = build_dataset(is_train=True, config=config) + dataset_train, config.MODEL.NUM_CLASSES = build_dataset( + is_train=True, config=config + ) config.freeze() dataset_val, _ = build_dataset(is_train=False, config=config) @@ -31,7 +33,8 @@ def build_loader_finetune(config): ) data_loader_train = DataLoader( - dataset_train, sampler=sampler_train, + dataset_train, + sampler=sampler_train, batch_size=config.DATA.BATCH_SIZE, num_workers=config.DATA.NUM_WORKERS, pin_memory=config.DATA.PIN_MEMORY, @@ -39,7 +42,8 @@ def build_loader_finetune(config): ) data_loader_val = DataLoader( - dataset_val, sampler=sampler_val, + dataset_val, + sampler=sampler_val, batch_size=config.DATA.BATCH_SIZE, num_workers=config.DATA.NUM_WORKERS, pin_memory=config.DATA.PIN_MEMORY, @@ -48,21 +52,31 @@ def build_loader_finetune(config): # 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 + mixup_active = ( + config.AUG.MIXUP > 0 + or config.AUG.CUTMIX > 0.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) + 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, data_loader_train, data_loader_val, mixup_fn def build_dataset(is_train, config): transform = build_transform(is_train, config) - - if config.DATA.DATASET == 'imagenet': - prefix = 'train' if is_train else 'val' + + if config.DATA.DATASET == "imagenet": + prefix = "train" if is_train else "val" root = os.path.join(config.DATA.DATA_PATH, prefix) dataset = datasets.ImageFolder(root, transform=transform) nb_classes = 1000 @@ -79,8 +93,12 @@ def build_transform(is_train, config): 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, + 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, @@ -89,7 +107,9 @@ def build_transform(is_train, config): if not resize_im: # replace RandomResizedCropAndInterpolation with # RandomCrop - transform.transforms[0] = transforms.RandomCrop(config.DATA.IMG_SIZE, padding=4) + transform.transforms[0] = transforms.RandomCrop( + config.DATA.IMG_SIZE, padding=4 + ) return transform t = [] @@ -97,14 +117,18 @@ def build_transform(is_train, config): if config.TEST.CROP: size = int((256 / 224) * config.DATA.IMG_SIZE) t.append( - transforms.Resize(size, interpolation=_pil_interp(config.DATA.INTERPOLATION)), + 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)) else: t.append( - transforms.Resize((config.DATA.IMG_SIZE, config.DATA.IMG_SIZE), - interpolation=_pil_interp(config.DATA.INTERPOLATION)) + transforms.Resize( + (config.DATA.IMG_SIZE, config.DATA.IMG_SIZE), + interpolation=_pil_interp(config.DATA.INTERPOLATION), + ) ) t.append(transforms.ToTensor()) diff --git a/data/data_simmim_pt.py b/data/data_simmim_pt.py index 0f2503a7..9826ce90 100644 --- a/data/data_simmim_pt.py +++ b/data/data_simmim_pt.py @@ -7,70 +7,81 @@ import math import random -import numpy as np +import numpy as np import torch import torch.distributed as dist import torchvision.transforms as T +from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data._utils.collate import default_collate from torchvision.datasets import ImageFolder -from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD class MaskGenerator: - def __init__(self, input_size=192, mask_patch_size=32, model_patch_size=4, mask_ratio=0.6): + def __init__( + self, input_size=192, mask_patch_size=32, model_patch_size=4, mask_ratio=0.6 + ): self.input_size = input_size self.mask_patch_size = mask_patch_size self.model_patch_size = model_patch_size self.mask_ratio = mask_ratio - + assert self.input_size % self.mask_patch_size == 0 assert self.mask_patch_size % self.model_patch_size == 0 - + self.rand_size = self.input_size // self.mask_patch_size self.scale = self.mask_patch_size // self.model_patch_size - - self.token_count = self.rand_size ** 2 + + self.token_count = self.rand_size**2 self.mask_count = int(np.ceil(self.token_count * self.mask_ratio)) - + def __call__(self): - mask_idx = np.random.permutation(self.token_count)[:self.mask_count] + mask_idx = np.random.permutation(self.token_count)[: self.mask_count] mask = np.zeros(self.token_count, dtype=int) mask[mask_idx] = 1 - + mask = mask.reshape((self.rand_size, self.rand_size)) mask = mask.repeat(self.scale, axis=0).repeat(self.scale, axis=1) - + return mask class SimMIMTransform: def __init__(self, config): - self.transform_img = T.Compose([ - T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img), - T.RandomResizedCrop(config.DATA.IMG_SIZE, scale=(0.67, 1.), ratio=(3. / 4., 4. / 3.)), - T.RandomHorizontalFlip(), - T.ToTensor(), - T.Normalize(mean=torch.tensor(IMAGENET_DEFAULT_MEAN),std=torch.tensor(IMAGENET_DEFAULT_STD)), - ]) - - if config.MODEL.TYPE in ['swin', 'swinv2']: - model_patch_size=config.MODEL.SWIN.PATCH_SIZE + self.transform_img = T.Compose( + [ + T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img), + T.RandomResizedCrop( + config.DATA.IMG_SIZE, + scale=(0.67, 1.0), + ratio=(3.0 / 4.0, 4.0 / 3.0), + ), + T.RandomHorizontalFlip(), + T.ToTensor(), + T.Normalize( + mean=torch.tensor(IMAGENET_DEFAULT_MEAN), + std=torch.tensor(IMAGENET_DEFAULT_STD), + ), + ] + ) + + if config.MODEL.TYPE in ["swin", "swinv2"]: + model_patch_size = config.MODEL.SWIN.PATCH_SIZE else: raise NotImplementedError - + self.mask_generator = MaskGenerator( input_size=config.DATA.IMG_SIZE, mask_patch_size=config.DATA.MASK_PATCH_SIZE, model_patch_size=model_patch_size, mask_ratio=config.DATA.MASK_RATIO, ) - + def __call__(self, img): img = self.transform_img(img) mask = self.mask_generator() - + return img, mask @@ -84,7 +95,9 @@ def collate_fn(batch): if batch[0][0][item_idx] is None: ret.append(None) else: - ret.append(default_collate([batch[i][0][item_idx] for i in range(batch_num)])) + ret.append( + default_collate([batch[i][0][item_idx] for i in range(batch_num)]) + ) ret.append(default_collate([batch[i][1] for i in range(batch_num)])) return ret @@ -92,8 +105,18 @@ def collate_fn(batch): def build_loader_simmim(config): transform = SimMIMTransform(config) dataset = ImageFolder(config.DATA.DATA_PATH, transform) - - sampler = DistributedSampler(dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank(), shuffle=True) - dataloader = DataLoader(dataset, config.DATA.BATCH_SIZE, sampler=sampler, num_workers=config.DATA.NUM_WORKERS, pin_memory=True, drop_last=True, collate_fn=collate_fn) - - return dataloader \ No newline at end of file + + sampler = DistributedSampler( + dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank(), shuffle=True + ) + dataloader = DataLoader( + dataset, + config.DATA.BATCH_SIZE, + sampler=sampler, + num_workers=config.DATA.NUM_WORKERS, + pin_memory=True, + drop_last=True, + collate_fn=collate_fn, + ) + + return dataloader diff --git a/data/imagenet22k_dataset.py b/data/imagenet22k_dataset.py index 5758060b..dc3210fe 100644 --- a/data/imagenet22k_dataset.py +++ b/data/imagenet22k_dataset.py @@ -1,16 +1,16 @@ -import os import json -import torch.utils.data as data +import os +import warnings + import numpy as np +import torch.utils.data as data from PIL import Image -import warnings - warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning) class IN22KDATASET(data.Dataset): - def __init__(self, root, ann_file='', transform=None, target_transform=None): + def __init__(self, root, ann_file="", transform=None, target_transform=None): super(IN22KDATASET, self).__init__() self.data_path = root @@ -40,7 +40,7 @@ def __getitem__(self, index): idb = self.database[index] # images - images = self._load_image(self.data_path + '/' + idb[0]).convert('RGB') + images = self._load_image(self.data_path + "/" + idb[0]).convert("RGB") if self.transform is not None: images = self.transform(images) diff --git a/data/zipreader.py b/data/zipreader.py index 060bc46a..1babf75e 100644 --- a/data/zipreader.py +++ b/data/zipreader.py @@ -5,23 +5,24 @@ # Written by Ze Liu # -------------------------------------------------------- +import io import os import zipfile -import io + import numpy as np -from PIL import Image -from PIL import ImageFile +from PIL import Image, ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True def is_zip_path(img_or_path): """judge if this is a zip path""" - return '.zip@' in img_or_path + return ".zip@" in img_or_path class ZipReader(object): """A class to read zipped files""" + zip_bank = dict() def __init__(self): @@ -31,18 +32,20 @@ def __init__(self): def get_zipfile(path): zip_bank = ZipReader.zip_bank if path not in zip_bank: - zfile = zipfile.ZipFile(path, 'r') + zfile = zipfile.ZipFile(path, "r") zip_bank[path] = zfile return zip_bank[path] @staticmethod def split_zip_style_path(path): - pos_at = path.index('@') - assert pos_at != -1, "character '@' is not found from the given path '%s'" % path - - zip_path = path[0: pos_at] - folder_path = path[pos_at + 1:] - folder_path = str.strip(folder_path, '/') + pos_at = path.index("@") + assert pos_at != -1, ( + "character '@' is not found from the given path '%s'" % path + ) + + zip_path = path[0:pos_at] + folder_path = path[pos_at + 1 :] + folder_path = str.strip(folder_path, "/") return zip_path, folder_path @staticmethod @@ -52,33 +55,37 @@ def list_folder(path): zfile = ZipReader.get_zipfile(zip_path) folder_list = [] for file_foler_name in zfile.namelist(): - file_foler_name = str.strip(file_foler_name, '/') - if file_foler_name.startswith(folder_path) and \ - len(os.path.splitext(file_foler_name)[-1]) == 0 and \ - file_foler_name != folder_path: + file_foler_name = str.strip(file_foler_name, "/") + if ( + file_foler_name.startswith(folder_path) + and len(os.path.splitext(file_foler_name)[-1]) == 0 + and file_foler_name != folder_path + ): if len(folder_path) == 0: folder_list.append(file_foler_name) else: - folder_list.append(file_foler_name[len(folder_path) + 1:]) + folder_list.append(file_foler_name[len(folder_path) + 1 :]) return folder_list @staticmethod def list_files(path, extension=None): if extension is None: - extension = ['.*'] + extension = [".*"] zip_path, folder_path = ZipReader.split_zip_style_path(path) zfile = ZipReader.get_zipfile(zip_path) file_lists = [] for file_foler_name in zfile.namelist(): - file_foler_name = str.strip(file_foler_name, '/') - if file_foler_name.startswith(folder_path) and \ - str.lower(os.path.splitext(file_foler_name)[-1]) in extension: + file_foler_name = str.strip(file_foler_name, "/") + if ( + file_foler_name.startswith(folder_path) + and str.lower(os.path.splitext(file_foler_name)[-1]) in extension + ): if len(folder_path) == 0: file_lists.append(file_foler_name) else: - file_lists.append(file_foler_name[len(folder_path) + 1:]) + file_lists.append(file_foler_name[len(folder_path) + 1 :]) return file_lists diff --git a/kernels/window_process/setup.py b/kernels/window_process/setup.py index c78526d0..c6cb5ea6 100644 --- a/kernels/window_process/setup.py +++ b/kernels/window_process/setup.py @@ -1,12 +1,16 @@ from setuptools import setup from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -setup(name='swin_window_process', +setup( + name="swin_window_process", ext_modules=[ - CUDAExtension('swin_window_process', [ - 'swin_window_process.cpp', - 'swin_window_process_kernel.cu', - ]) + CUDAExtension( + "swin_window_process", + [ + "swin_window_process.cpp", + "swin_window_process_kernel.cu", + ], + ) ], - cmdclass={'build_ext': BuildExtension}) \ No newline at end of file + cmdclass={"build_ext": BuildExtension}, +) diff --git a/kernels/window_process/unit_test.py b/kernels/window_process/unit_test.py index 65dee566..743d7691 100644 --- a/kernels/window_process/unit_test.py +++ b/kernels/window_process/unit_test.py @@ -4,22 +4,25 @@ # Licensed under The MIT License [see LICENSE for details] # -------------------------------------------------------- -import torch -import swin_window_process import random import time import unittest +import swin_window_process +import torch + class WindowProcess(torch.autograd.Function): @staticmethod def forward(ctx, input, B, H, W, C, shift_size, window_size): - output = swin_window_process.roll_and_window_partition_forward(input, B, H, W, C, shift_size, window_size) + output = swin_window_process.roll_and_window_partition_forward( + input, B, H, W, C, shift_size, window_size + ) ctx.B = B ctx.H = H - ctx.W = W - ctx.C = C + ctx.W = W + ctx.C = C ctx.shift_size = shift_size ctx.window_size = window_size return output @@ -28,24 +31,28 @@ def forward(ctx, input, B, H, W, C, shift_size, window_size): def backward(ctx, grad_in): B = ctx.B H = ctx.H - W = ctx.W - C = ctx.C + W = ctx.W + C = ctx.C shift_size = ctx.shift_size window_size = ctx.window_size - grad_out = swin_window_process.roll_and_window_partition_backward(grad_in, B, H, W, C, shift_size, window_size) + grad_out = swin_window_process.roll_and_window_partition_backward( + grad_in, B, H, W, C, shift_size, window_size + ) return grad_out, None, None, None, None, None, None, None class WindowProcessReverse(torch.autograd.Function): @staticmethod def forward(ctx, input, B, H, W, C, shift_size, window_size): - output = swin_window_process.window_merge_and_roll_forward(input, B, H, W, C, shift_size, window_size) + output = swin_window_process.window_merge_and_roll_forward( + input, B, H, W, C, shift_size, window_size + ) ctx.B = B ctx.H = H - ctx.W = W - ctx.C = C + ctx.W = W + ctx.C = C ctx.shift_size = shift_size ctx.window_size = window_size @@ -55,12 +62,14 @@ def forward(ctx, input, B, H, W, C, shift_size, window_size): def backward(ctx, grad_in): B = ctx.B H = ctx.H - W = ctx.W - C = ctx.C + W = ctx.W + C = ctx.C shift_size = ctx.shift_size window_size = ctx.window_size - grad_out = swin_window_process.window_merge_and_roll_backward(grad_in, B, H, W, C, shift_size, window_size) + grad_out = swin_window_process.window_merge_and_roll_backward( + grad_in, B, H, W, C, shift_size, window_size + ) return grad_out, None, None, None, None, None, None, None @@ -74,9 +83,12 @@ def window_partition(x, window_size): """ B, H, W, C = x.shape x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) - windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + windows = ( + x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + ) return windows + def window_reverse(windows, window_size, H, W): """ Args: @@ -88,7 +100,9 @@ def window_reverse(windows, window_size, H, W): x: (B, H, W, C) """ B = int(windows.shape[0] / (H * W / window_size / window_size)) - x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = windows.view( + B, H // window_size, W // window_size, window_size, window_size, -1 + ) x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) return x @@ -119,6 +133,7 @@ def copy_one_tensor(input, requires_grad=True): input1 = input.clone().detach().requires_grad_(requires_grad).cuda() return input1 + class Test_WindowProcess(unittest.TestCase): def setUp(self): self.B = 192 @@ -129,10 +144,12 @@ def setUp(self): self.window_size = 7 self.nH = self.H // self.window_size self.nW = self.W // self.window_size - + def test_roll_and_window_partition_forward(self, dtype=torch.float32): - input = torch.randn((self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True).cuda() - + input = torch.randn( + (self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True + ).cuda() + input1 = copy_one_tensor(input, True) input2 = copy_one_tensor(input, True) @@ -140,15 +157,28 @@ def test_roll_and_window_partition_forward(self, dtype=torch.float32): # ori expected = pyt_forward(input1, self.shift_size, self.window_size) # fused kernel - fused_output = WindowProcess.apply(input2, self.B, self.H, self.W, self.C, -self.shift_size, self.window_size) - + fused_output = WindowProcess.apply( + input2, + self.B, + self.H, + self.W, + self.C, + -self.shift_size, + self.window_size, + ) + self.assertTrue(torch.equal(expected, fused_output)) - #self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) - + # self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) + def test_roll_and_window_partition_backward(self, dtype=torch.float32): - input = torch.randn((self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True).cuda() - d_loss_tensor = torch.randn((self.B*self.nW*self.nH, self.window_size, self.window_size, self.C), dtype=dtype).cuda() - + input = torch.randn( + (self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True + ).cuda() + d_loss_tensor = torch.randn( + (self.B * self.nW * self.nH, self.window_size, self.window_size, self.C), + dtype=dtype, + ).cuda() + input1 = copy_one_tensor(input, True) input2 = copy_one_tensor(input, True) @@ -156,64 +186,105 @@ def test_roll_and_window_partition_backward(self, dtype=torch.float32): expected = pyt_forward(input1, self.shift_size, self.window_size) expected.backward(d_loss_tensor) # fused kernel - fused_output = WindowProcess.apply(input2, self.B, self.H, self.W, self.C, -self.shift_size, self.window_size) + fused_output = WindowProcess.apply( + input2, self.B, self.H, self.W, self.C, -self.shift_size, self.window_size + ) fused_output.backward(d_loss_tensor) - + self.assertTrue(torch.equal(expected, fused_output)) - #self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) + # self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) def test_window_merge_and_roll_forward(self, dtype=torch.float32): - input = torch.randn((self.B*self.nH*self.nW, self.window_size, self.window_size, self.C), dtype=dtype, requires_grad=True).cuda() - + input = torch.randn( + (self.B * self.nH * self.nW, self.window_size, self.window_size, self.C), + dtype=dtype, + requires_grad=True, + ).cuda() + input1 = copy_one_tensor(input, True) input2 = copy_one_tensor(input, True) with torch.no_grad(): # ori - expected = reverse_pyt_forward(input1, self.shift_size, self.window_size, self.H, self.W) + expected = reverse_pyt_forward( + input1, self.shift_size, self.window_size, self.H, self.W + ) # fused kernel - fused_output = WindowProcessReverse.apply(input2, self.B, self.H, self.W, self.C, self.shift_size, self.window_size) - + fused_output = WindowProcessReverse.apply( + input2, + self.B, + self.H, + self.W, + self.C, + self.shift_size, + self.window_size, + ) + self.assertTrue(torch.equal(expected, fused_output)) - #self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) - + # self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) def test_window_merge_and_roll_backward(self, dtype=torch.float32): - input = torch.randn((self.B*self.nH*self.nW, self.window_size, self.window_size, self.C), dtype=dtype, requires_grad=True).cuda() - d_loss_tensor = torch.randn((self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True).cuda() - + input = torch.randn( + (self.B * self.nH * self.nW, self.window_size, self.window_size, self.C), + dtype=dtype, + requires_grad=True, + ).cuda() + d_loss_tensor = torch.randn( + (self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True + ).cuda() + input1 = copy_one_tensor(input, True) input2 = copy_one_tensor(input, True) # ori - expected = reverse_pyt_forward(input1, self.shift_size, self.window_size, self.H, self.W) + expected = reverse_pyt_forward( + input1, self.shift_size, self.window_size, self.H, self.W + ) expected.backward(d_loss_tensor) # fused kernel - fused_output = WindowProcessReverse.apply(input2, self.B, self.H, self.W, self.C, self.shift_size, self.window_size) + fused_output = WindowProcessReverse.apply( + input2, self.B, self.H, self.W, self.C, self.shift_size, self.window_size + ) fused_output.backward(d_loss_tensor) - + self.assertTrue(torch.equal(expected, fused_output)) - #self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) + # self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) def test_forward_backward_speed(self, dtype=torch.float32, times=1000): - input = torch.randn((self.B*self.nH*self.nW, self.window_size, self.window_size, self.C), dtype=dtype, requires_grad=True).cuda() - d_loss_tensor = torch.randn((self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True).cuda() - + input = torch.randn( + (self.B * self.nH * self.nW, self.window_size, self.window_size, self.C), + dtype=dtype, + requires_grad=True, + ).cuda() + d_loss_tensor = torch.randn( + (self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True + ).cuda() + input1 = copy_one_tensor(input, True) input2 = copy_one_tensor(input, True) # SwinTransformer official def run_pyt(t=1000): for _ in range(t): - expected = reverse_pyt_forward(input1, self.shift_size, self.window_size, self.H, self.W) + expected = reverse_pyt_forward( + input1, self.shift_size, self.window_size, self.H, self.W + ) expected.backward(d_loss_tensor) # my op def run_fusedop(t=1000): for _ in range(t): - fused_output = WindowProcessReverse.apply(input2, self.B, self.H, self.W, self.C, self.shift_size, self.window_size) + fused_output = WindowProcessReverse.apply( + input2, + self.B, + self.H, + self.W, + self.C, + self.shift_size, + self.window_size, + ) fused_output.backward(d_loss_tensor) - + torch.cuda.synchronize() t1 = time.time() run_pyt(t=times) @@ -224,10 +295,10 @@ def run_fusedop(t=1000): t3 = time.time() self.assertTrue((t3 - t2) < (t2 - t1)) - print('Run {} times'.format(times)) - print('Original time cost: {}'.format(t2 - t1)) - print('Fused op time cost: {}'.format(t3 - t2)) - + print("Run {} times".format(times)) + print("Original time cost: {}".format(t2 - t1)) + print("Fused op time cost: {}".format(t3 - t2)) + def test_roll_and_window_partition_forward_fp16(self, dtype=torch.float16): self.test_roll_and_window_partition_forward(dtype=dtype) @@ -236,7 +307,7 @@ def test_roll_and_window_partition_backward_fp16(self, dtype=torch.float16): def test_window_merge_and_roll_forward_fp16(self, dtype=torch.float16): self.test_window_merge_and_roll_forward(dtype=dtype) - + def test_window_merge_and_roll_backward_fp16(self, dtype=torch.float16): self.test_window_merge_and_roll_backward(dtype=dtype) @@ -244,7 +315,7 @@ def test_forward_backward_speed_fp16(self, dtype=torch.float16, times=1000): self.test_forward_backward_speed(dtype=dtype, times=times) -if __name__ == '__main__': - print('Pass only two tensors are exactly the same (using torch.equal).\n') +if __name__ == "__main__": + print("Pass only two tensors are exactly the same (using torch.equal).\n") torch.manual_seed(0) unittest.main(verbosity=2) diff --git a/kernels/window_process/window_process.py b/kernels/window_process/window_process.py index ee43e9e9..d482d4f7 100644 --- a/kernels/window_process/window_process.py +++ b/kernels/window_process/window_process.py @@ -4,19 +4,21 @@ # Licensed under The MIT License [see LICENSE for details] # -------------------------------------------------------- -import torch import swin_window_process +import torch class WindowProcess(torch.autograd.Function): @staticmethod def forward(ctx, input, B, H, W, C, shift_size, window_size): - output = swin_window_process.roll_and_window_partition_forward(input, B, H, W, C, shift_size, window_size) + output = swin_window_process.roll_and_window_partition_forward( + input, B, H, W, C, shift_size, window_size + ) ctx.B = B ctx.H = H - ctx.W = W - ctx.C = C + ctx.W = W + ctx.C = C ctx.shift_size = shift_size ctx.window_size = window_size return output @@ -25,24 +27,28 @@ def forward(ctx, input, B, H, W, C, shift_size, window_size): def backward(ctx, grad_in): B = ctx.B H = ctx.H - W = ctx.W - C = ctx.C + W = ctx.W + C = ctx.C shift_size = ctx.shift_size window_size = ctx.window_size - grad_out = swin_window_process.roll_and_window_partition_backward(grad_in, B, H, W, C, shift_size, window_size) + grad_out = swin_window_process.roll_and_window_partition_backward( + grad_in, B, H, W, C, shift_size, window_size + ) return grad_out, None, None, None, None, None, None, None class WindowProcessReverse(torch.autograd.Function): @staticmethod def forward(ctx, input, B, H, W, C, shift_size, window_size): - output = swin_window_process.window_merge_and_roll_forward(input, B, H, W, C, shift_size, window_size) + output = swin_window_process.window_merge_and_roll_forward( + input, B, H, W, C, shift_size, window_size + ) ctx.B = B ctx.H = H - ctx.W = W - ctx.C = C + ctx.W = W + ctx.C = C ctx.shift_size = shift_size ctx.window_size = window_size @@ -52,12 +58,14 @@ def forward(ctx, input, B, H, W, C, shift_size, window_size): def backward(ctx, grad_in): B = ctx.B H = ctx.H - W = ctx.W - C = ctx.C + W = ctx.W + C = ctx.C shift_size = ctx.shift_size window_size = ctx.window_size - #grad_out = ctx.saved_tensors[0] - #grad_out = torch.zeros((B, H, W, C), dtype=dtype).cuda() - grad_out = swin_window_process.window_merge_and_roll_backward(grad_in, B, H, W, C, shift_size, window_size) + # grad_out = ctx.saved_tensors[0] + # grad_out = torch.zeros((B, H, W, C), dtype=dtype).cuda() + grad_out = swin_window_process.window_merge_and_roll_backward( + grad_in, B, H, W, C, shift_size, window_size + ) return grad_out, None, None, None, None, None, None, None diff --git a/logger.py b/logger.py index a066e55b..d2e0b4b1 100644 --- a/logger.py +++ b/logger.py @@ -5,37 +5,44 @@ # Written by Ze Liu # -------------------------------------------------------- +import functools +import logging import os import sys -import logging -import functools + from termcolor import colored @functools.lru_cache() -def create_logger(output_dir, dist_rank=0, name=''): +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' + 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')) + 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 = 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')) + file_handler.setFormatter(logging.Formatter(fmt=fmt, datefmt="%Y-%m-%d %H:%M:%S")) logger.addHandler(file_handler) return logger diff --git a/lr_scheduler.py b/lr_scheduler.py index a2122e5d..e991c529 100644 --- a/lr_scheduler.py +++ b/lr_scheduler.py @@ -9,8 +9,8 @@ import torch from timm.scheduler.cosine_lr import CosineLRScheduler -from timm.scheduler.step_lr import StepLRScheduler from timm.scheduler.scheduler import Scheduler +from timm.scheduler.step_lr import StepLRScheduler def build_scheduler(config, optimizer, n_iter_per_epoch): @@ -20,11 +20,13 @@ def build_scheduler(config, optimizer, n_iter_per_epoch): multi_steps = [i * n_iter_per_epoch for i in config.TRAIN.LR_SCHEDULER.MULTISTEPS] lr_scheduler = None - if config.TRAIN.LR_SCHEDULER.NAME == 'cosine': + if config.TRAIN.LR_SCHEDULER.NAME == "cosine": lr_scheduler = CosineLRScheduler( optimizer, - t_initial=(num_steps - warmup_steps) if config.TRAIN.LR_SCHEDULER.WARMUP_PREFIX else num_steps, - t_mul=1., + t_initial=(num_steps - warmup_steps) + if config.TRAIN.LR_SCHEDULER.WARMUP_PREFIX + else num_steps, + t_mul=1.0, lr_min=config.TRAIN.MIN_LR, warmup_lr_init=config.TRAIN.WARMUP_LR, warmup_t=warmup_steps, @@ -32,7 +34,7 @@ def build_scheduler(config, optimizer, n_iter_per_epoch): t_in_epochs=False, warmup_prefix=config.TRAIN.LR_SCHEDULER.WARMUP_PREFIX, ) - elif config.TRAIN.LR_SCHEDULER.NAME == 'linear': + elif config.TRAIN.LR_SCHEDULER.NAME == "linear": lr_scheduler = LinearLRScheduler( optimizer, t_initial=num_steps, @@ -41,7 +43,7 @@ def build_scheduler(config, optimizer, n_iter_per_epoch): warmup_t=warmup_steps, t_in_epochs=False, ) - elif config.TRAIN.LR_SCHEDULER.NAME == 'step': + elif config.TRAIN.LR_SCHEDULER.NAME == "step": lr_scheduler = StepLRScheduler( optimizer, decay_t=decay_steps, @@ -50,7 +52,7 @@ def build_scheduler(config, optimizer, n_iter_per_epoch): warmup_t=warmup_steps, t_in_epochs=False, ) - elif config.TRAIN.LR_SCHEDULER.NAME == 'multistep': + elif config.TRAIN.LR_SCHEDULER.NAME == "multistep": lr_scheduler = MultiStepLRScheduler( optimizer, milestones=multi_steps, @@ -64,23 +66,29 @@ def build_scheduler(config, optimizer, n_iter_per_epoch): class LinearLRScheduler(Scheduler): - def __init__(self, - optimizer: torch.optim.Optimizer, - t_initial: int, - lr_min_rate: float, - warmup_t=0, - warmup_lr_init=0., - t_in_epochs=True, - noise_range_t=None, - noise_pct=0.67, - noise_std=1.0, - noise_seed=42, - initialize=True, - ) -> None: + def __init__( + self, + optimizer: torch.optim.Optimizer, + t_initial: int, + lr_min_rate: float, + warmup_t=0, + warmup_lr_init=0.0, + t_in_epochs=True, + noise_range_t=None, + noise_pct=0.67, + noise_std=1.0, + noise_seed=42, + initialize=True, + ) -> None: super().__init__( - optimizer, param_group_field="lr", - noise_range_t=noise_range_t, noise_pct=noise_pct, noise_std=noise_std, noise_seed=noise_seed, - initialize=initialize) + optimizer, + param_group_field="lr", + noise_range_t=noise_range_t, + noise_pct=noise_pct, + noise_std=noise_std, + noise_seed=noise_seed, + initialize=initialize, + ) self.t_initial = t_initial self.lr_min_rate = lr_min_rate @@ -88,7 +96,9 @@ def __init__(self, self.warmup_lr_init = warmup_lr_init self.t_in_epochs = t_in_epochs if self.warmup_t: - self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in self.base_values] + self.warmup_steps = [ + (v - warmup_lr_init) / self.warmup_t for v in self.base_values + ] super().update_groups(self.warmup_lr_init) else: self.warmup_steps = [1 for _ in self.base_values] @@ -99,7 +109,10 @@ def _get_lr(self, t): else: t = t - self.warmup_t total_t = self.t_initial - self.warmup_t - lrs = [v - ((v - v * self.lr_min_rate) * (t / total_t)) for v in self.base_values] + lrs = [ + v - ((v - v * self.lr_min_rate) * (t / total_t)) + for v in self.base_values + ] return lrs def get_epoch_values(self, epoch: int): @@ -116,27 +129,40 @@ def get_update_values(self, num_updates: int): class MultiStepLRScheduler(Scheduler): - def __init__(self, optimizer: torch.optim.Optimizer, milestones, gamma=0.1, warmup_t=0, warmup_lr_init=0, t_in_epochs=True) -> None: + def __init__( + self, + optimizer: torch.optim.Optimizer, + milestones, + gamma=0.1, + warmup_t=0, + warmup_lr_init=0, + t_in_epochs=True, + ) -> None: super().__init__(optimizer, param_group_field="lr") - + self.milestones = milestones self.gamma = gamma self.warmup_t = warmup_t self.warmup_lr_init = warmup_lr_init self.t_in_epochs = t_in_epochs if self.warmup_t: - self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in self.base_values] + self.warmup_steps = [ + (v - warmup_lr_init) / self.warmup_t for v in self.base_values + ] super().update_groups(self.warmup_lr_init) else: self.warmup_steps = [1 for _ in self.base_values] - + assert self.warmup_t <= min(self.milestones) - + def _get_lr(self, t): if t < self.warmup_t: lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps] else: - lrs = [v * (self.gamma ** bisect.bisect_right(self.milestones, t)) for v in self.base_values] + lrs = [ + v * (self.gamma ** bisect.bisect_right(self.milestones, t)) + for v in self.base_values + ] return lrs def get_epoch_values(self, epoch: int): diff --git a/main.py b/main.py index 84230ea7..cfe333fc 100644 --- a/main.py +++ b/main.py @@ -5,74 +5,129 @@ # Written by Ze Liu # -------------------------------------------------------- -import os -import time -import json -import random import argparse import datetime -import numpy as np +import json +import os +import random +import time +import numpy as np import torch import torch.backends.cudnn as cudnn import torch.distributed as dist - from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy -from timm.utils import accuracy, AverageMeter +from timm.utils import AverageMeter, accuracy from config import get_config -from models import build_model from data import build_loader +from logger import create_logger from lr_scheduler import build_scheduler +from models import build_model from optimizer import build_optimizer -from logger import create_logger -from utils import load_checkpoint, load_pretrained, save_checkpoint, NativeScalerWithGradNormCount, auto_resume_helper, \ - reduce_tensor +from utils import ( + NativeScalerWithGradNormCount, + auto_resume_helper, + load_checkpoint, + load_pretrained, + reduce_tensor, + save_checkpoint, +) def parse_option(): - parser = argparse.ArgumentParser('Swin Transformer training and evaluation script', add_help=False) - parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file', ) + parser = argparse.ArgumentParser( + "Swin Transformer training and evaluation script", add_help=False + ) + parser.add_argument( + "--cfg", + type=str, + required=True, + metavar="FILE", + help="path to config file", + ) parser.add_argument( "--opts", help="Modify config options by adding 'KEY VALUE' pairs. ", default=None, - nargs='+', + nargs="+", ) # easy config modification - parser.add_argument('--batch-size', type=int, help="batch size for single GPU") - parser.add_argument('--data-path', type=str, help='path to dataset') - parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset') - parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'], - help='no: no cache, ' - 'full: cache all data, ' - 'part: sharding the dataset into nonoverlapping pieces and only cache one piece') - parser.add_argument('--pretrained', - help='pretrained weight from checkpoint, could be imagenet22k pretrained weight') - parser.add_argument('--resume', help='resume from checkpoint') - parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps") - parser.add_argument('--use-checkpoint', action='store_true', - help="whether to use gradient checkpointing to save memory") - parser.add_argument('--disable_amp', action='store_true', help='Disable pytorch amp') - parser.add_argument('--amp-opt-level', type=str, choices=['O0', 'O1', 'O2'], - help='mixed precision opt level, if O0, no amp is used (deprecated!)') - parser.add_argument('--output', default='output', type=str, metavar='PATH', - help='root of output folder, the full path is // (default: output)') - parser.add_argument('--tag', help='tag of experiment') - parser.add_argument('--eval', action='store_true', help='Perform evaluation only') - parser.add_argument('--throughput', action='store_true', help='Test throughput only') + parser.add_argument("--batch-size", type=int, help="batch size for single GPU") + parser.add_argument("--data-path", type=str, help="path to dataset") + parser.add_argument( + "--zip", + action="store_true", + help="use zipped dataset instead of folder dataset", + ) + parser.add_argument( + "--cache-mode", + type=str, + default="part", + choices=["no", "full", "part"], + help="no: no cache, " + "full: cache all data, " + "part: sharding the dataset into nonoverlapping pieces and only cache one piece", + ) + parser.add_argument( + "--pretrained", + help="pretrained weight from checkpoint, could be imagenet22k pretrained weight", + ) + parser.add_argument("--resume", help="resume from checkpoint") + parser.add_argument( + "--accumulation-steps", type=int, help="gradient accumulation steps" + ) + parser.add_argument( + "--use-checkpoint", + action="store_true", + help="whether to use gradient checkpointing to save memory", + ) + parser.add_argument( + "--disable_amp", action="store_true", help="Disable pytorch amp" + ) + parser.add_argument( + "--amp-opt-level", + type=str, + choices=["O0", "O1", "O2"], + help="mixed precision opt level, if O0, no amp is used (deprecated!)", + ) + parser.add_argument( + "--output", + default="output", + type=str, + metavar="PATH", + help="root of output folder, the full path is // (default: output)", + ) + parser.add_argument("--tag", help="tag of experiment") + parser.add_argument("--eval", action="store_true", help="Perform evaluation only") + parser.add_argument( + "--throughput", action="store_true", help="Test throughput only" + ) # distributed training - parser.add_argument("--local_rank", type=int, required=True, help='local rank for DistributedDataParallel') + parser.add_argument( + "--local_rank", + type=int, + required=True, + help="local rank for DistributedDataParallel", + ) # for acceleration - parser.add_argument('--fused_window_process', action='store_true', - help='Fused window shift & window partition, similar for reversed part.') - parser.add_argument('--fused_layernorm', action='store_true', help='Use fused layernorm.') + parser.add_argument( + "--fused_window_process", + action="store_true", + help="Fused window shift & window partition, similar for reversed part.", + ) + parser.add_argument( + "--fused_layernorm", action="store_true", help="Use fused layernorm." + ) ## overwrite optimizer in config (*.yaml) if specified, e.g., fused_adam/fused_lamb - parser.add_argument('--optim', type=str, - help='overwrite optimizer if provided, can be adamw/sgd/fused_adam/fused_lamb.') + parser.add_argument( + "--optim", + type=str, + help="overwrite optimizer if provided, can be adamw/sgd/fused_adam/fused_lamb.", + ) args, unparsed = parser.parse_known_args() @@ -82,7 +137,13 @@ def parse_option(): def main(config): - dataset_train, dataset_val, data_loader_train, data_loader_val, mixup_fn = build_loader(config) + ( + dataset_train, + dataset_val, + data_loader_train, + data_loader_val, + mixup_fn, + ) = build_loader(config) logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}") model = build_model(config) @@ -90,7 +151,7 @@ def main(config): n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) logger.info(f"number of params: {n_parameters}") - if hasattr(model, 'flops'): + if hasattr(model, "flops"): flops = model.flops() logger.info(f"number of GFLOPs: {flops / 1e9}") @@ -98,18 +159,22 @@ def main(config): model_without_ddp = model optimizer = build_optimizer(config, model) - model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False) + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False + ) loss_scaler = NativeScalerWithGradNormCount() if config.TRAIN.ACCUMULATION_STEPS > 1: - lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train) // config.TRAIN.ACCUMULATION_STEPS) + lr_scheduler = build_scheduler( + config, optimizer, len(data_loader_train) // config.TRAIN.ACCUMULATION_STEPS + ) else: lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train)) - if config.AUG.MIXUP > 0.: + if config.AUG.MIXUP > 0.0: # smoothing is handled with mixup label transform criterion = SoftTargetCrossEntropy() - elif config.MODEL.LABEL_SMOOTHING > 0.: + elif config.MODEL.LABEL_SMOOTHING > 0.0: criterion = LabelSmoothingCrossEntropy(smoothing=config.MODEL.LABEL_SMOOTHING) else: criterion = torch.nn.CrossEntropyLoss() @@ -120,25 +185,33 @@ def main(config): resume_file = auto_resume_helper(config.OUTPUT) if resume_file: if config.MODEL.RESUME: - logger.warning(f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}") + logger.warning( + f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}" + ) config.defrost() config.MODEL.RESUME = resume_file config.freeze() - logger.info(f'auto resuming from {resume_file}') + logger.info(f"auto resuming from {resume_file}") else: - logger.info(f'no checkpoint found in {config.OUTPUT}, ignoring auto resume') + logger.info(f"no checkpoint found in {config.OUTPUT}, ignoring auto resume") if config.MODEL.RESUME: - max_accuracy = load_checkpoint(config, model_without_ddp, optimizer, lr_scheduler, loss_scaler, logger) + max_accuracy = load_checkpoint( + config, model_without_ddp, optimizer, lr_scheduler, loss_scaler, logger + ) acc1, acc5, loss = validate(config, data_loader_val, model) - logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") + logger.info( + f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" + ) if config.EVAL_MODE: return if config.MODEL.PRETRAINED and (not config.MODEL.RESUME): load_pretrained(config, model_without_ddp, logger) acc1, acc5, loss = validate(config, data_loader_val, model) - logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") + logger.info( + f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" + ) if config.THROUGHPUT_MODE: throughput(data_loader_val, model, logger) @@ -149,23 +222,54 @@ def main(config): for epoch in range(config.TRAIN.START_EPOCH, config.TRAIN.EPOCHS): data_loader_train.sampler.set_epoch(epoch) - train_one_epoch(config, model, criterion, data_loader_train, optimizer, epoch, mixup_fn, lr_scheduler, - loss_scaler) - if dist.get_rank() == 0 and (epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1)): - save_checkpoint(config, epoch, model_without_ddp, max_accuracy, optimizer, lr_scheduler, loss_scaler, - logger) + train_one_epoch( + config, + model, + criterion, + data_loader_train, + optimizer, + epoch, + mixup_fn, + lr_scheduler, + loss_scaler, + ) + if dist.get_rank() == 0 and ( + epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1) + ): + save_checkpoint( + config, + epoch, + model_without_ddp, + max_accuracy, + optimizer, + lr_scheduler, + loss_scaler, + logger, + ) acc1, acc5, loss = validate(config, data_loader_val, model) - logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") + logger.info( + f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" + ) max_accuracy = max(max_accuracy, acc1) - logger.info(f'Max accuracy: {max_accuracy:.2f}%') + logger.info(f"Max accuracy: {max_accuracy:.2f}%") total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) - logger.info('Training time {}'.format(total_time_str)) - - -def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mixup_fn, lr_scheduler, loss_scaler): + logger.info("Training time {}".format(total_time_str)) + + +def train_one_epoch( + config, + model, + criterion, + data_loader, + optimizer, + epoch, + mixup_fn, + lr_scheduler, + loss_scaler, +): model.train() optimizer.zero_grad() @@ -190,13 +294,22 @@ def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mix loss = loss / config.TRAIN.ACCUMULATION_STEPS # this attribute is added by timm on one optimizer (adahessian) - is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order - grad_norm = loss_scaler(loss, optimizer, clip_grad=config.TRAIN.CLIP_GRAD, - parameters=model.parameters(), create_graph=is_second_order, - update_grad=(idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0) + is_second_order = ( + hasattr(optimizer, "is_second_order") and optimizer.is_second_order + ) + grad_norm = loss_scaler( + loss, + optimizer, + clip_grad=config.TRAIN.CLIP_GRAD, + parameters=model.parameters(), + create_graph=is_second_order, + update_grad=(idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0, + ) if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0: optimizer.zero_grad() - lr_scheduler.step_update((epoch * num_steps + idx) // config.TRAIN.ACCUMULATION_STEPS) + lr_scheduler.step_update( + (epoch * num_steps + idx) // config.TRAIN.ACCUMULATION_STEPS + ) loss_scale_value = loss_scaler.state_dict()["scale"] torch.cuda.synchronize() @@ -209,20 +322,23 @@ def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mix end = time.time() if idx % config.PRINT_FREQ == 0: - lr = optimizer.param_groups[0]['lr'] - wd = optimizer.param_groups[0]['weight_decay'] + lr = optimizer.param_groups[0]["lr"] + wd = optimizer.param_groups[0]["weight_decay"] memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) etas = batch_time.avg * (num_steps - idx) logger.info( - f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t' - f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t wd {wd:.4f}\t' - f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t' - f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t' - f'grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t' - f'loss_scale {scaler_meter.val:.4f} ({scaler_meter.avg:.4f})\t' - f'mem {memory_used:.0f}MB') + f"Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t" + f"eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t wd {wd:.4f}\t" + f"time {batch_time.val:.4f} ({batch_time.avg:.4f})\t" + f"loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t" + f"grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t" + f"loss_scale {scaler_meter.val:.4f} ({scaler_meter.avg:.4f})\t" + f"mem {memory_used:.0f}MB" + ) epoch_time = time.time() - start - logger.info(f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}") + logger.info( + f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}" + ) @torch.no_grad() @@ -263,13 +379,14 @@ def validate(config, data_loader, model): if idx % config.PRINT_FREQ == 0: memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) logger.info( - f'Test: [{idx}/{len(data_loader)}]\t' - f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' - f'Loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t' - f'Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t' - f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t' - f'Mem {memory_used:.0f}MB') - logger.info(f' * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}') + f"Test: [{idx}/{len(data_loader)}]\t" + f"Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t" + f"Loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t" + f"Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t" + f"Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t" + f"Mem {memory_used:.0f}MB" + ) + logger.info(f" * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}") return acc1_meter.avg, acc5_meter.avg, loss_meter.avg @@ -289,25 +406,29 @@ def throughput(data_loader, model, logger): model(images) torch.cuda.synchronize() tic2 = time.time() - logger.info(f"batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}") + logger.info( + f"batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}" + ) return -if __name__ == '__main__': +if __name__ == "__main__": args, config = parse_option() if config.AMP_OPT_LEVEL: print("[warning] Apex amp has been deprecated, please use pytorch amp instead!") - if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: rank = int(os.environ["RANK"]) - world_size = int(os.environ['WORLD_SIZE']) + world_size = int(os.environ["WORLD_SIZE"]) print(f"RANK and WORLD_SIZE in environ: {rank}/{world_size}") else: rank = -1 world_size = -1 torch.cuda.set_device(config.LOCAL_RANK) - torch.distributed.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank) + torch.distributed.init_process_group( + backend="nccl", init_method="env://", world_size=world_size, rank=rank + ) torch.distributed.barrier() seed = config.SEED + dist.get_rank() @@ -318,13 +439,21 @@ def throughput(data_loader, model, logger): cudnn.benchmark = True # linear scale the learning rate according to total batch size, may not be optimal - linear_scaled_lr = config.TRAIN.BASE_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 - linear_scaled_warmup_lr = config.TRAIN.WARMUP_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 - linear_scaled_min_lr = config.TRAIN.MIN_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + linear_scaled_lr = ( + config.TRAIN.BASE_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) + linear_scaled_warmup_lr = ( + config.TRAIN.WARMUP_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) + linear_scaled_min_lr = ( + config.TRAIN.MIN_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) # gradient accumulation also need to scale the learning rate if config.TRAIN.ACCUMULATION_STEPS > 1: linear_scaled_lr = linear_scaled_lr * config.TRAIN.ACCUMULATION_STEPS - linear_scaled_warmup_lr = linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS + linear_scaled_warmup_lr = ( + linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS + ) linear_scaled_min_lr = linear_scaled_min_lr * config.TRAIN.ACCUMULATION_STEPS config.defrost() config.TRAIN.BASE_LR = linear_scaled_lr @@ -333,7 +462,9 @@ def throughput(data_loader, model, logger): config.freeze() os.makedirs(config.OUTPUT, exist_ok=True) - logger = create_logger(output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}") + logger = create_logger( + output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" + ) if dist.get_rank() == 0: path = os.path.join(config.OUTPUT, "config.json") diff --git a/main_moe.py b/main_moe.py index acf5d205..082d9c59 100644 --- a/main_moe.py +++ b/main_moe.py @@ -5,70 +5,117 @@ # Written by Ze Liu # -------------------------------------------------------- -from tutel import system - -import os -import time -import json -import random import argparse import datetime -import numpy as np +import json +import os +import random +import time from functools import partial + +import numpy as np import torch import torch.backends.cudnn as cudnn import torch.distributed as dist - from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy -from timm.utils import accuracy, AverageMeter +from timm.utils import AverageMeter, accuracy +from tutel import system from config import get_config -from models import build_model from data import build_loader +from logger import create_logger from lr_scheduler import build_scheduler +from models import build_model from optimizer import build_optimizer -from logger import create_logger from utils import NativeScalerWithGradNormCount, reduce_tensor -from utils_moe import load_checkpoint, load_pretrained, save_checkpoint, auto_resume_helper, hook_scale_grad +from utils_moe import ( + auto_resume_helper, + hook_scale_grad, + load_checkpoint, + load_pretrained, + save_checkpoint, +) -assert torch.__version__ >= '1.8.0', "DDP-based MoE requires Pytorch >= 1.8.0" +assert torch.__version__ >= "1.8.0", "DDP-based MoE requires Pytorch >= 1.8.0" def parse_option(): - parser = argparse.ArgumentParser('Swin Transformer training and evaluation script', add_help=False) - parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file', ) + parser = argparse.ArgumentParser( + "Swin Transformer training and evaluation script", add_help=False + ) + parser.add_argument( + "--cfg", + type=str, + required=True, + metavar="FILE", + help="path to config file", + ) parser.add_argument( "--opts", help="Modify config options by adding 'KEY VALUE' pairs. ", default=None, - nargs='+', + nargs="+", ) # easy config modification - parser.add_argument('--batch-size', type=int, help="batch size for single GPU") - parser.add_argument('--data-path', type=str, help='path to dataset') - parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset') - parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'], - help='no: no cache, ' - 'full: cache all data, ' - 'part: sharding the dataset into nonoverlapping pieces and only cache one piece') - parser.add_argument('--pretrained', - help='pretrained weight from checkpoint, could be imagenet22k pretrained weight') - parser.add_argument('--resume', help='resume from checkpoint') - parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps") - parser.add_argument('--use-checkpoint', action='store_true', - help="whether to use gradient checkpointing to save memory") - parser.add_argument('--disable_amp', action='store_true', help='Disable pytorch amp') - parser.add_argument('--amp-opt-level', type=str, choices=['O0', 'O1', 'O2'], - help='mixed precision opt level, if O0, no amp is used (deprecated!)') - parser.add_argument('--output', default='output', type=str, metavar='PATH', - help='root of output folder, the full path is // (default: output)') - parser.add_argument('--tag', help='tag of experiment') - parser.add_argument('--eval', action='store_true', help='Perform evaluation only') - parser.add_argument('--throughput', action='store_true', help='Test throughput only') + parser.add_argument("--batch-size", type=int, help="batch size for single GPU") + parser.add_argument("--data-path", type=str, help="path to dataset") + parser.add_argument( + "--zip", + action="store_true", + help="use zipped dataset instead of folder dataset", + ) + parser.add_argument( + "--cache-mode", + type=str, + default="part", + choices=["no", "full", "part"], + help="no: no cache, " + "full: cache all data, " + "part: sharding the dataset into nonoverlapping pieces and only cache one piece", + ) + parser.add_argument( + "--pretrained", + help="pretrained weight from checkpoint, could be imagenet22k pretrained weight", + ) + parser.add_argument("--resume", help="resume from checkpoint") + parser.add_argument( + "--accumulation-steps", type=int, help="gradient accumulation steps" + ) + parser.add_argument( + "--use-checkpoint", + action="store_true", + help="whether to use gradient checkpointing to save memory", + ) + parser.add_argument( + "--disable_amp", action="store_true", help="Disable pytorch amp" + ) + parser.add_argument( + "--amp-opt-level", + type=str, + choices=["O0", "O1", "O2"], + help="mixed precision opt level, if O0, no amp is used (deprecated!)", + ) + parser.add_argument( + "--output", + default="output", + type=str, + metavar="PATH", + help="root of output folder, the full path is // (default: output)", + ) + parser.add_argument("--tag", help="tag of experiment") + parser.add_argument("--eval", action="store_true", help="Perform evaluation only") + parser.add_argument( + "--throughput", action="store_true", help="Test throughput only" + ) # distributed training - parser.add_argument("--local_rank", type=int, required=True, help='local rank for DistributedDataParallel') + parser.add_argument( + "--local_rank", + type=int, + required=True, + help="local rank for DistributedDataParallel", + ) args, unparsed = parser.parse_known_args() @@ -78,7 +125,13 @@ def parse_option(): def main(config): - dataset_train, dataset_val, data_loader_train, data_loader_val, mixup_fn = build_loader(config) + ( + dataset_train, + dataset_val, + data_loader_train, + data_loader_val, + mixup_fn, + ) = build_loader(config) logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}") model = build_model(config) @@ -86,18 +139,32 @@ def main(config): # For Tutel MoE for name, param in model.named_parameters(): - if param.requires_grad == True and hasattr(param, 'skip_allreduce') and param.skip_allreduce is True: + if ( + param.requires_grad == True + and hasattr(param, "skip_allreduce") + and param.skip_allreduce is True + ): model.add_param_to_skip_allreduce(name) param.register_hook(partial(hook_scale_grad, dist.get_world_size())) - logger.info(f"[rank{dist.get_rank()}] [{name}] skip all_reduce and div {dist.get_world_size()} for grad") + logger.info( + f"[rank{dist.get_rank()}] [{name}] skip all_reduce and div {dist.get_world_size()} for grad" + ) - n_parameters_single = sum(p.numel() * model.sharded_count if hasattr(p, 'skip_allreduce') - else p.numel() for p in model.parameters() if p.requires_grad) + n_parameters_single = sum( + p.numel() * model.sharded_count if hasattr(p, "skip_allreduce") else p.numel() + for p in model.parameters() + if p.requires_grad + ) logger.info(f"number of params single: {n_parameters_single}") - n_parameters_whole = sum(p.numel() * model.sharded_count * model.global_experts if hasattr(p, 'skip_allreduce') - else p.numel() for p in model.parameters() if p.requires_grad) + n_parameters_whole = sum( + p.numel() * model.sharded_count * model.global_experts + if hasattr(p, "skip_allreduce") + else p.numel() + for p in model.parameters() + if p.requires_grad + ) logger.info(f"number of params whole: {n_parameters_whole}") - if hasattr(model, 'flops'): + if hasattr(model, "flops"): flops = model.flops() logger.info(f"number of GFLOPs: {flops / 1e9}") @@ -105,18 +172,22 @@ def main(config): model_without_ddp = model optimizer = build_optimizer(config, model) - model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False) + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False + ) loss_scaler = NativeScalerWithGradNormCount() if config.TRAIN.ACCUMULATION_STEPS > 1: - lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train) // config.TRAIN.ACCUMULATION_STEPS) + lr_scheduler = build_scheduler( + config, optimizer, len(data_loader_train) // config.TRAIN.ACCUMULATION_STEPS + ) else: lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train)) - if config.AUG.MIXUP > 0.: + if config.AUG.MIXUP > 0.0: # smoothing is handled with mixup label transform criterion = SoftTargetCrossEntropy() - elif config.MODEL.LABEL_SMOOTHING > 0.: + elif config.MODEL.LABEL_SMOOTHING > 0.0: criterion = LabelSmoothingCrossEntropy(smoothing=config.MODEL.LABEL_SMOOTHING) else: criterion = torch.nn.CrossEntropyLoss() @@ -127,25 +198,33 @@ def main(config): resume_file = auto_resume_helper(config.OUTPUT, config.TRAIN.MOE.SAVE_MASTER) if resume_file: if config.MODEL.RESUME: - logger.warning(f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}") + logger.warning( + f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}" + ) config.defrost() config.MODEL.RESUME = resume_file config.freeze() - logger.info(f'auto resuming from {resume_file}') + logger.info(f"auto resuming from {resume_file}") else: - logger.info(f'no checkpoint found in {config.OUTPUT}, ignoring auto resume') + logger.info(f"no checkpoint found in {config.OUTPUT}, ignoring auto resume") if config.MODEL.RESUME: - max_accuracy = load_checkpoint(config, model_without_ddp, optimizer, lr_scheduler, loss_scaler, logger) + max_accuracy = load_checkpoint( + config, model_without_ddp, optimizer, lr_scheduler, loss_scaler, logger + ) acc1, acc5, loss = validate(config, data_loader_val, model) - logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") + logger.info( + f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" + ) if config.EVAL_MODE: return if config.MODEL.PRETRAINED and (not config.MODEL.RESUME): load_pretrained(config, model_without_ddp, logger) acc1, acc5, loss = validate(config, data_loader_val, model) - logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") + logger.info( + f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" + ) if config.EVAL_MODE: return @@ -158,24 +237,62 @@ def main(config): for epoch in range(config.TRAIN.START_EPOCH, config.TRAIN.EPOCHS): data_loader_train.sampler.set_epoch(epoch) - train_one_epoch(config, model, criterion, data_loader_train, optimizer, epoch, mixup_fn, lr_scheduler, - loss_scaler) - if (epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1)): - save_checkpoint(config, epoch, model_without_ddp, max_accuracy, optimizer, lr_scheduler, loss_scaler, - logger) + train_one_epoch( + config, + model, + criterion, + data_loader_train, + optimizer, + epoch, + mixup_fn, + lr_scheduler, + loss_scaler, + ) + if epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1): + save_checkpoint( + config, + epoch, + model_without_ddp, + max_accuracy, + optimizer, + lr_scheduler, + loss_scaler, + logger, + ) acc1, acc5, loss = validate(config, data_loader_val, model) - logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") + logger.info( + f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" + ) max_accuracy = max(max_accuracy, acc1) - logger.info(f'Max accuracy: {max_accuracy:.2f}%') - save_checkpoint(config, 'final', model_without_ddp, max_accuracy, optimizer, lr_scheduler, loss_scaler, - logger, zero_redundancy=True) + logger.info(f"Max accuracy: {max_accuracy:.2f}%") + save_checkpoint( + config, + "final", + model_without_ddp, + max_accuracy, + optimizer, + lr_scheduler, + loss_scaler, + logger, + zero_redundancy=True, + ) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) - logger.info('Training time {}'.format(total_time_str)) - - -def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mixup_fn, lr_scheduler, loss_scaler): + logger.info("Training time {}".format(total_time_str)) + + +def train_one_epoch( + config, + model, + criterion, + data_loader, + optimizer, + epoch, + mixup_fn, + lr_scheduler, + loss_scaler, +): model.train() optimizer.zero_grad() @@ -203,20 +320,31 @@ def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mix loss = loss / config.TRAIN.ACCUMULATION_STEPS # this attribute is added by timm on one optimizer (adahessian) - is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order - grad_norm = loss_scaler(loss, optimizer, clip_grad=config.TRAIN.CLIP_GRAD, - parameters=model.parameters(), create_graph=is_second_order, - update_grad=(idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0) + is_second_order = ( + hasattr(optimizer, "is_second_order") and optimizer.is_second_order + ) + grad_norm = loss_scaler( + loss, + optimizer, + clip_grad=config.TRAIN.CLIP_GRAD, + parameters=model.parameters(), + create_graph=is_second_order, + update_grad=(idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0, + ) if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0: optimizer.zero_grad() - lr_scheduler.step_update((epoch * num_steps + idx) // config.TRAIN.ACCUMULATION_STEPS) + lr_scheduler.step_update( + (epoch * num_steps + idx) // config.TRAIN.ACCUMULATION_STEPS + ) loss_scale_value = loss_scaler.state_dict()["scale"] torch.cuda.synchronize() loss_meter.update(loss.item(), targets.size(0)) loss_cls_meter.update(l_cls.item(), targets.size(0)) - loss_aux_meter.update(l_aux if isinstance(l_aux, float) else l_aux.item(), targets.size(0)) + loss_aux_meter.update( + l_aux if isinstance(l_aux, float) else l_aux.item(), targets.size(0) + ) if grad_norm is not None: # loss_scaler return None if not update norm_meter.update(grad_norm) scaler_meter.update(loss_scale_value) @@ -224,22 +352,25 @@ def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mix end = time.time() if idx % config.PRINT_FREQ == 0: - lr = optimizer.param_groups[0]['lr'] - wd = optimizer.param_groups[0]['weight_decay'] + lr = optimizer.param_groups[0]["lr"] + wd = optimizer.param_groups[0]["weight_decay"] memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) etas = batch_time.avg * (num_steps - idx) logger.info( - f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t' - f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t wd {wd:.4f}\t' - f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t' - f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t' - f'loss-cls {loss_cls_meter.val:.4f} ({loss_cls_meter.avg:.4f})\t' - f'loss-aux {loss_aux_meter.val:.4f} ({loss_aux_meter.avg:.4f})\t' - f'grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t' - f'loss_scale {scaler_meter.val:.4f} ({scaler_meter.avg:.4f})\t' - f'mem {memory_used:.0f}MB') + f"Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t" + f"eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t wd {wd:.4f}\t" + f"time {batch_time.val:.4f} ({batch_time.avg:.4f})\t" + f"loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t" + f"loss-cls {loss_cls_meter.val:.4f} ({loss_cls_meter.avg:.4f})\t" + f"loss-aux {loss_aux_meter.val:.4f} ({loss_aux_meter.avg:.4f})\t" + f"grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t" + f"loss_scale {scaler_meter.val:.4f} ({scaler_meter.avg:.4f})\t" + f"mem {memory_used:.0f}MB" + ) epoch_time = time.time() - start - logger.info(f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}") + logger.info( + f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}" + ) @torch.no_grad() @@ -270,7 +401,9 @@ def validate(config, data_loader, model): acc5 = reduce_tensor(acc5) loss_cls_meter.update(l_cls.item(), target.size(0)) - loss_aux_meter.update(l_aux if isinstance(l_aux, float) else l_aux.item(), target.size(0)) + loss_aux_meter.update( + l_aux if isinstance(l_aux, float) else l_aux.item(), target.size(0) + ) acc1_meter.update(acc1.item(), target.size(0)) acc5_meter.update(acc5.item(), target.size(0)) @@ -281,14 +414,15 @@ def validate(config, data_loader, model): if idx % config.PRINT_FREQ == 0: memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) logger.info( - f'Test: [{idx}/{len(data_loader)}]\t' - f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' - f'Loss-Cls {loss_cls_meter.val:.4f} ({loss_cls_meter.avg:.4f})\t' - f'Loss-Aux {loss_aux_meter.val:.4f} ({loss_aux_meter.avg:.4f})\t' - f'Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t' - f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t' - f'Mem {memory_used:.0f}MB') - logger.info(f' * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}') + f"Test: [{idx}/{len(data_loader)}]\t" + f"Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t" + f"Loss-Cls {loss_cls_meter.val:.4f} ({loss_cls_meter.avg:.4f})\t" + f"Loss-Aux {loss_aux_meter.val:.4f} ({loss_aux_meter.avg:.4f})\t" + f"Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t" + f"Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t" + f"Mem {memory_used:.0f}MB" + ) + logger.info(f" * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}") return acc1_meter.avg, acc5_meter.avg, loss_cls_meter.avg @@ -308,25 +442,29 @@ def throughput(data_loader, model, logger): model(images) torch.cuda.synchronize() tic2 = time.time() - logger.info(f"batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}") + logger.info( + f"batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}" + ) return -if __name__ == '__main__': +if __name__ == "__main__": args, config = parse_option() if config.AMP_OPT_LEVEL: print("[warning] Apex amp has been deprecated, please use pytorch amp instead!") - if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: rank = int(os.environ["RANK"]) - world_size = int(os.environ['WORLD_SIZE']) + world_size = int(os.environ["WORLD_SIZE"]) print(f"RANK and WORLD_SIZE in environ: {rank}/{world_size}") else: rank = -1 world_size = -1 torch.cuda.set_device(config.LOCAL_RANK) - torch.distributed.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank) + torch.distributed.init_process_group( + backend="nccl", init_method="env://", world_size=world_size, rank=rank + ) torch.distributed.barrier() seed = config.SEED + dist.get_rank() @@ -337,13 +475,21 @@ def throughput(data_loader, model, logger): cudnn.benchmark = True # linear scale the learning rate according to total batch size, may not be optimal - linear_scaled_lr = config.TRAIN.BASE_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 - linear_scaled_warmup_lr = config.TRAIN.WARMUP_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 - linear_scaled_min_lr = config.TRAIN.MIN_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + linear_scaled_lr = ( + config.TRAIN.BASE_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) + linear_scaled_warmup_lr = ( + config.TRAIN.WARMUP_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) + linear_scaled_min_lr = ( + config.TRAIN.MIN_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) # gradient accumulation also need to scale the learning rate if config.TRAIN.ACCUMULATION_STEPS > 1: linear_scaled_lr = linear_scaled_lr * config.TRAIN.ACCUMULATION_STEPS - linear_scaled_warmup_lr = linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS + linear_scaled_warmup_lr = ( + linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS + ) linear_scaled_min_lr = linear_scaled_min_lr * config.TRAIN.ACCUMULATION_STEPS config.defrost() config.TRAIN.BASE_LR = linear_scaled_lr @@ -352,7 +498,9 @@ def throughput(data_loader, model, logger): config.freeze() os.makedirs(config.OUTPUT, exist_ok=True) - logger = create_logger(output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}") + logger = create_logger( + output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" + ) if dist.get_rank() == 0: path = os.path.join(config.OUTPUT, "config.json") diff --git a/main_simmim_ft.py b/main_simmim_ft.py index 067dfbb0..2f0df1d7 100644 --- a/main_simmim_ft.py +++ b/main_simmim_ft.py @@ -6,58 +6,87 @@ # Modified by Zhenda Xie # -------------------------------------------------------- -import os -import time import argparse import datetime -import numpy as np +import os +import time +import numpy as np import torch import torch.backends.cudnn as cudnn -import torch.distributed as dist import torch.cuda.amp as amp - +import torch.distributed as dist from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy -from timm.utils import accuracy, AverageMeter +from timm.utils import AverageMeter, accuracy from config import get_config -from models import build_model from data import build_loader +from logger import create_logger from lr_scheduler import build_scheduler +from models import build_model from optimizer import build_optimizer -from logger import create_logger -from utils_simmim import load_checkpoint, load_pretrained, save_checkpoint, get_grad_norm, auto_resume_helper, reduce_tensor +from utils_simmim import ( + auto_resume_helper, + get_grad_norm, + load_checkpoint, + load_pretrained, + reduce_tensor, + save_checkpoint, +) def parse_option(): - parser = argparse.ArgumentParser('SimMIM fine-tuning script', add_help=False) - parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file', ) + parser = argparse.ArgumentParser("SimMIM fine-tuning script", add_help=False) + parser.add_argument( + "--cfg", + type=str, + required=True, + metavar="FILE", + help="path to config file", + ) parser.add_argument( "--opts", help="Modify config options by adding 'KEY VALUE' pairs. ", default=None, - nargs='+', + nargs="+", ) # easy config modification - parser.add_argument('--batch-size', type=int, help="batch size for single GPU") - parser.add_argument('--data-path', type=str, help='path to dataset') - parser.add_argument('--pretrained', type=str, help='path to pre-trained model') - parser.add_argument('--resume', help='resume from checkpoint') - parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps") - parser.add_argument('--use-checkpoint', action='store_true', - help="whether to use gradient checkpointing to save memory") - parser.add_argument('--enable-amp', action='store_true') - parser.add_argument('--disable-amp', action='store_false', dest='enable_amp') + parser.add_argument("--batch-size", type=int, help="batch size for single GPU") + parser.add_argument("--data-path", type=str, help="path to dataset") + parser.add_argument("--pretrained", type=str, help="path to pre-trained model") + parser.add_argument("--resume", help="resume from checkpoint") + parser.add_argument( + "--accumulation-steps", type=int, help="gradient accumulation steps" + ) + parser.add_argument( + "--use-checkpoint", + action="store_true", + help="whether to use gradient checkpointing to save memory", + ) + parser.add_argument("--enable-amp", action="store_true") + parser.add_argument("--disable-amp", action="store_false", dest="enable_amp") parser.set_defaults(enable_amp=True) - parser.add_argument('--output', default='output', type=str, metavar='PATH', - help='root of output folder, the full path is // (default: output)') - parser.add_argument('--tag', help='tag of experiment') - parser.add_argument('--eval', action='store_true', help='Perform evaluation only') - parser.add_argument('--throughput', action='store_true', help='Test throughput only') + parser.add_argument( + "--output", + default="output", + type=str, + metavar="PATH", + help="root of output folder, the full path is // (default: output)", + ) + parser.add_argument("--tag", help="tag of experiment") + parser.add_argument("--eval", action="store_true", help="Perform evaluation only") + parser.add_argument( + "--throughput", action="store_true", help="Test throughput only" + ) # distributed training - parser.add_argument("--local_rank", type=int, required=True, help='local rank for DistributedDataParallel') + parser.add_argument( + "--local_rank", + type=int, + required=True, + help="local rank for DistributedDataParallel", + ) args = parser.parse_args() @@ -67,7 +96,13 @@ def parse_option(): def main(config): - dataset_train, dataset_val, data_loader_train, data_loader_val, mixup_fn = build_loader(config, simmim=True, is_pretrain=False) + ( + dataset_train, + dataset_val, + data_loader_train, + data_loader_val, + mixup_fn, + ) = build_loader(config, simmim=True, is_pretrain=False) logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}") model = build_model(config, is_pretrain=False) @@ -75,22 +110,24 @@ def main(config): logger.info(str(model)) optimizer = build_optimizer(config, model, simmim=True, is_pretrain=False) - model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False) + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False + ) model_without_ddp = model.module n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) logger.info(f"number of params: {n_parameters}") - if hasattr(model_without_ddp, 'flops'): + if hasattr(model_without_ddp, "flops"): flops = model_without_ddp.flops() logger.info(f"number of GFLOPs: {flops / 1e9}") lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train)) scaler = amp.GradScaler() - if config.AUG.MIXUP > 0.: + if config.AUG.MIXUP > 0.0: # smoothing is handled with mixup label transform criterion = SoftTargetCrossEntropy() - elif config.MODEL.LABEL_SMOOTHING > 0.: + elif config.MODEL.LABEL_SMOOTHING > 0.0: criterion = LabelSmoothingCrossEntropy(smoothing=config.MODEL.LABEL_SMOOTHING) else: criterion = torch.nn.CrossEntropyLoss() @@ -101,25 +138,33 @@ def main(config): resume_file = auto_resume_helper(config.OUTPUT, logger) if resume_file: if config.MODEL.RESUME: - logger.warning(f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}") + logger.warning( + f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}" + ) config.defrost() config.MODEL.RESUME = resume_file config.freeze() - logger.info(f'auto resuming from {resume_file}') + logger.info(f"auto resuming from {resume_file}") else: - logger.info(f'no checkpoint found in {config.OUTPUT}, ignoring auto resume') + logger.info(f"no checkpoint found in {config.OUTPUT}, ignoring auto resume") if config.MODEL.RESUME: - max_accuracy = load_checkpoint(config, model_without_ddp, optimizer, lr_scheduler, scaler, logger) + max_accuracy = load_checkpoint( + config, model_without_ddp, optimizer, lr_scheduler, scaler, logger + ) acc1, acc5, loss = validate(config, data_loader_val, model) - logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") + logger.info( + f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" + ) if config.EVAL_MODE: return if config.MODEL.PRETRAINED and (not config.MODEL.RESUME): load_pretrained(config, model_without_ddp, logger) acc1, acc5, loss = validate(config, data_loader_val, model) - logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") + logger.info( + f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" + ) if config.THROUGHPUT_MODE: throughput(data_loader_val, model, logger) @@ -130,25 +175,60 @@ def main(config): for epoch in range(config.TRAIN.START_EPOCH, config.TRAIN.EPOCHS): data_loader_train.sampler.set_epoch(epoch) - train_one_epoch(config, model, criterion, data_loader_train, optimizer, epoch, mixup_fn, lr_scheduler, scaler) - if dist.get_rank() == 0 and (epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1)): - save_checkpoint(config, epoch, model_without_ddp, max_accuracy, optimizer, lr_scheduler, scaler, logger) + train_one_epoch( + config, + model, + criterion, + data_loader_train, + optimizer, + epoch, + mixup_fn, + lr_scheduler, + scaler, + ) + if dist.get_rank() == 0 and ( + epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1) + ): + save_checkpoint( + config, + epoch, + model_without_ddp, + max_accuracy, + optimizer, + lr_scheduler, + scaler, + logger, + ) acc1, acc5, loss = validate(config, data_loader_val, model) - logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") + logger.info( + f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" + ) max_accuracy = max(max_accuracy, acc1) - logger.info(f'Max accuracy: {max_accuracy:.2f}%') + logger.info(f"Max accuracy: {max_accuracy:.2f}%") total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) - logger.info('Training time {}'.format(total_time_str)) - - -def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mixup_fn, lr_scheduler, scaler): + logger.info("Training time {}".format(total_time_str)) + + +def train_one_epoch( + config, + model, + criterion, + data_loader, + optimizer, + epoch, + mixup_fn, + lr_scheduler, + scaler, +): model.train() optimizer.zero_grad() - - logger.info(f'Current learning rate for different parameter groups: {[it["lr"] for it in optimizer.param_groups]}') + + logger.info( + f'Current learning rate for different parameter groups: {[it["lr"] for it in optimizer.param_groups]}' + ) num_steps = len(data_loader) batch_time = AverageMeter() @@ -173,7 +253,9 @@ def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mix scaler.scale(loss).backward() if config.TRAIN.CLIP_GRAD: scaler.unscale_(optimizer) - grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.TRAIN.CLIP_GRAD) + grad_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), config.TRAIN.CLIP_GRAD + ) else: grad_norm = get_grad_norm(model.parameters()) if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0: @@ -187,7 +269,9 @@ def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mix scaler.scale(loss).backward() if config.TRAIN.CLIP_GRAD: scaler.unscale_(optimizer) - grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.TRAIN.CLIP_GRAD) + grad_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), config.TRAIN.CLIP_GRAD + ) else: grad_norm = get_grad_norm(model.parameters()) scaler.step(optimizer) @@ -203,19 +287,22 @@ def train_one_epoch(config, model, criterion, data_loader, optimizer, epoch, mix end = time.time() if idx % config.PRINT_FREQ == 0: - lr = optimizer.param_groups[-1]['lr'] + lr = optimizer.param_groups[-1]["lr"] memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) etas = batch_time.avg * (num_steps - idx) logger.info( - f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t' - f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t' - f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t' - f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t' - f'grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t' - f'loss_scale {loss_scale_meter.val:.4f} ({loss_scale_meter.avg:.4f})\t' - f'mem {memory_used:.0f}MB') + f"Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t" + f"eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t" + f"time {batch_time.val:.4f} ({batch_time.avg:.4f})\t" + f"loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t" + f"grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t" + f"loss_scale {loss_scale_meter.val:.4f} ({loss_scale_meter.avg:.4f})\t" + f"mem {memory_used:.0f}MB" + ) epoch_time = time.time() - start - logger.info(f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}") + logger.info( + f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}" + ) @torch.no_grad() @@ -255,13 +342,14 @@ def validate(config, data_loader, model): if idx % config.PRINT_FREQ == 0: memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) logger.info( - f'Test: [{idx}/{len(data_loader)}]\t' - f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' - f'Loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t' - f'Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t' - f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t' - f'Mem {memory_used:.0f}MB') - logger.info(f' * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}') + f"Test: [{idx}/{len(data_loader)}]\t" + f"Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t" + f"Loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t" + f"Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t" + f"Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t" + f"Mem {memory_used:.0f}MB" + ) + logger.info(f" * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}") return acc1_meter.avg, acc5_meter.avg, loss_meter.avg @@ -281,22 +369,26 @@ def throughput(data_loader, model, logger): model(images) torch.cuda.synchronize() tic2 = time.time() - logger.info(f"batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}") + logger.info( + f"batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}" + ) return -if __name__ == '__main__': +if __name__ == "__main__": _, config = parse_option() - if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: rank = int(os.environ["RANK"]) - world_size = int(os.environ['WORLD_SIZE']) + world_size = int(os.environ["WORLD_SIZE"]) print(f"RANK and WORLD_SIZE in environ: {rank}/{world_size}") else: rank = -1 world_size = -1 torch.cuda.set_device(config.LOCAL_RANK) - torch.distributed.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank) + torch.distributed.init_process_group( + backend="nccl", init_method="env://", world_size=world_size, rank=rank + ) torch.distributed.barrier() seed = config.SEED + dist.get_rank() @@ -305,13 +397,21 @@ def throughput(data_loader, model, logger): cudnn.benchmark = True # linear scale the learning rate according to total batch size, may not be optimal - linear_scaled_lr = config.TRAIN.BASE_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 - linear_scaled_warmup_lr = config.TRAIN.WARMUP_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 - linear_scaled_min_lr = config.TRAIN.MIN_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + linear_scaled_lr = ( + config.TRAIN.BASE_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) + linear_scaled_warmup_lr = ( + config.TRAIN.WARMUP_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) + linear_scaled_min_lr = ( + config.TRAIN.MIN_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) # gradient accumulation also need to scale the learning rate if config.TRAIN.ACCUMULATION_STEPS > 1: linear_scaled_lr = linear_scaled_lr * config.TRAIN.ACCUMULATION_STEPS - linear_scaled_warmup_lr = linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS + linear_scaled_warmup_lr = ( + linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS + ) linear_scaled_min_lr = linear_scaled_min_lr * config.TRAIN.ACCUMULATION_STEPS config.defrost() config.TRAIN.BASE_LR = linear_scaled_lr @@ -320,7 +420,9 @@ def throughput(data_loader, model, logger): config.freeze() os.makedirs(config.OUTPUT, exist_ok=True) - logger = create_logger(output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}") + logger = create_logger( + output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" + ) if dist.get_rank() == 0: path = os.path.join(config.OUTPUT, "config.json") @@ -331,4 +433,4 @@ def throughput(data_loader, model, logger): # print config logger.info(config.dump()) - main(config) \ No newline at end of file + main(config) diff --git a/main_simmim_pt.py b/main_simmim_pt.py index 6591d214..b4d91dad 100644 --- a/main_simmim_pt.py +++ b/main_simmim_pt.py @@ -6,53 +6,79 @@ # Modified by Zhenda Xie # -------------------------------------------------------- -import os -import time import argparse import datetime -import numpy as np +import os +import time +import numpy as np import torch import torch.backends.cudnn as cudnn -import torch.distributed as dist import torch.cuda.amp as amp +import torch.distributed as dist from timm.utils import AverageMeter from config import get_config -from models import build_model from data import build_loader +from logger import create_logger from lr_scheduler import build_scheduler +from models import build_model from optimizer import build_optimizer -from logger import create_logger -from utils_simmim import load_checkpoint, save_checkpoint, get_grad_norm, auto_resume_helper +from utils_simmim import ( + auto_resume_helper, + get_grad_norm, + load_checkpoint, + save_checkpoint, +) def parse_option(): - parser = argparse.ArgumentParser('SimMIM pre-training script', add_help=False) - parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file', ) + parser = argparse.ArgumentParser("SimMIM pre-training script", add_help=False) + parser.add_argument( + "--cfg", + type=str, + required=True, + metavar="FILE", + help="path to config file", + ) parser.add_argument( "--opts", help="Modify config options by adding 'KEY VALUE' pairs. ", default=None, - nargs='+', + nargs="+", ) # easy config modification - parser.add_argument('--batch-size', type=int, help="batch size for single GPU") - parser.add_argument('--data-path', type=str, help='path to dataset') - parser.add_argument('--resume', help='resume from checkpoint') - parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps") - parser.add_argument('--use-checkpoint', action='store_true', - help="whether to use gradient checkpointing to save memory") - parser.add_argument('--enable-amp', action='store_true') - parser.add_argument('--disable-amp', action='store_false', dest='enable_amp') + parser.add_argument("--batch-size", type=int, help="batch size for single GPU") + parser.add_argument("--data-path", type=str, help="path to dataset") + parser.add_argument("--resume", help="resume from checkpoint") + parser.add_argument( + "--accumulation-steps", type=int, help="gradient accumulation steps" + ) + parser.add_argument( + "--use-checkpoint", + action="store_true", + help="whether to use gradient checkpointing to save memory", + ) + parser.add_argument("--enable-amp", action="store_true") + parser.add_argument("--disable-amp", action="store_false", dest="enable_amp") parser.set_defaults(enable_amp=True) - parser.add_argument('--output', default='output', type=str, metavar='PATH', - help='root of output folder, the full path is // (default: output)') - parser.add_argument('--tag', help='tag of experiment') + parser.add_argument( + "--output", + default="output", + type=str, + metavar="PATH", + help="root of output folder, the full path is // (default: output)", + ) + parser.add_argument("--tag", help="tag of experiment") # distributed training - parser.add_argument("--local_rank", type=int, required=True, help='local rank for DistributedDataParallel') + parser.add_argument( + "--local_rank", + type=int, + required=True, + help="local rank for DistributedDataParallel", + ) args = parser.parse_args() @@ -70,12 +96,14 @@ def main(config): logger.info(str(model)) optimizer = build_optimizer(config, model, simmim=True, is_pretrain=True) - model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False) + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False + ) model_without_ddp = model.module n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) logger.info(f"number of params: {n_parameters}") - if hasattr(model_without_ddp, 'flops'): + if hasattr(model_without_ddp, "flops"): flops = model_without_ddp.flops() logger.info(f"number of GFLOPs: {flops / 1e9}") @@ -86,29 +114,46 @@ def main(config): resume_file = auto_resume_helper(config.OUTPUT, logger) if resume_file: if config.MODEL.RESUME: - logger.warning(f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}") + logger.warning( + f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}" + ) config.defrost() config.MODEL.RESUME = resume_file config.freeze() - logger.info(f'auto resuming from {resume_file}') + logger.info(f"auto resuming from {resume_file}") else: - logger.info(f'no checkpoint found in {config.OUTPUT}, ignoring auto resume') + logger.info(f"no checkpoint found in {config.OUTPUT}, ignoring auto resume") if config.MODEL.RESUME: - load_checkpoint(config, model_without_ddp, optimizer, lr_scheduler, scaler, logger) + load_checkpoint( + config, model_without_ddp, optimizer, lr_scheduler, scaler, logger + ) logger.info("Start training") start_time = time.time() for epoch in range(config.TRAIN.START_EPOCH, config.TRAIN.EPOCHS): data_loader_train.sampler.set_epoch(epoch) - train_one_epoch(config, model, data_loader_train, optimizer, epoch, lr_scheduler, scaler) - if dist.get_rank() == 0 and (epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1)): - save_checkpoint(config, epoch, model_without_ddp, 0., optimizer, lr_scheduler, scaler, logger) + train_one_epoch( + config, model, data_loader_train, optimizer, epoch, lr_scheduler, scaler + ) + if dist.get_rank() == 0 and ( + epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1) + ): + save_checkpoint( + config, + epoch, + model_without_ddp, + 0.0, + optimizer, + lr_scheduler, + scaler, + logger, + ) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) - logger.info('Training time {}'.format(total_time_str)) + logger.info("Training time {}".format(total_time_str)) def train_one_epoch(config, model, data_loader, optimizer, epoch, lr_scheduler, scaler): @@ -135,7 +180,9 @@ def train_one_epoch(config, model, data_loader, optimizer, epoch, lr_scheduler, scaler.scale(loss).backward() if config.TRAIN.CLIP_GRAD: scaler.unscale_(optimizer) - grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.TRAIN.CLIP_GRAD) + grad_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), config.TRAIN.CLIP_GRAD + ) else: grad_norm = get_grad_norm(model.parameters()) if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0: @@ -148,7 +195,9 @@ def train_one_epoch(config, model, data_loader, optimizer, epoch, lr_scheduler, scaler.scale(loss).backward() if config.TRAIN.CLIP_GRAD: scaler.unscale_(optimizer) - grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.TRAIN.CLIP_GRAD) + grad_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), config.TRAIN.CLIP_GRAD + ) else: grad_norm = get_grad_norm(model.parameters()) scaler.step(optimizer) @@ -164,33 +213,38 @@ def train_one_epoch(config, model, data_loader, optimizer, epoch, lr_scheduler, end = time.time() if idx % config.PRINT_FREQ == 0: - lr = optimizer.param_groups[0]['lr'] + lr = optimizer.param_groups[0]["lr"] memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) etas = batch_time.avg * (num_steps - idx) logger.info( - f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t' - f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t' - f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t' - f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t' - f'grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t' - f'loss_scale {loss_scale_meter.val:.4f} ({loss_scale_meter.avg:.4f})\t' - f'mem {memory_used:.0f}MB') + f"Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t" + f"eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t" + f"time {batch_time.val:.4f} ({batch_time.avg:.4f})\t" + f"loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t" + f"grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t" + f"loss_scale {loss_scale_meter.val:.4f} ({loss_scale_meter.avg:.4f})\t" + f"mem {memory_used:.0f}MB" + ) epoch_time = time.time() - start - logger.info(f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}") + logger.info( + f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}" + ) -if __name__ == '__main__': +if __name__ == "__main__": _, config = parse_option() - if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: rank = int(os.environ["RANK"]) - world_size = int(os.environ['WORLD_SIZE']) + world_size = int(os.environ["WORLD_SIZE"]) print(f"RANK and WORLD_SIZE in environ: {rank}/{world_size}") else: rank = -1 world_size = -1 torch.cuda.set_device(config.LOCAL_RANK) - torch.distributed.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank) + torch.distributed.init_process_group( + backend="nccl", init_method="env://", world_size=world_size, rank=rank + ) torch.distributed.barrier() seed = config.SEED + dist.get_rank() @@ -199,13 +253,21 @@ def train_one_epoch(config, model, data_loader, optimizer, epoch, lr_scheduler, cudnn.benchmark = True # linear scale the learning rate according to total batch size, may not be optimal - linear_scaled_lr = config.TRAIN.BASE_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 - linear_scaled_warmup_lr = config.TRAIN.WARMUP_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 - linear_scaled_min_lr = config.TRAIN.MIN_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + linear_scaled_lr = ( + config.TRAIN.BASE_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) + linear_scaled_warmup_lr = ( + config.TRAIN.WARMUP_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) + linear_scaled_min_lr = ( + config.TRAIN.MIN_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + ) # gradient accumulation also need to scale the learning rate if config.TRAIN.ACCUMULATION_STEPS > 1: linear_scaled_lr = linear_scaled_lr * config.TRAIN.ACCUMULATION_STEPS - linear_scaled_warmup_lr = linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS + linear_scaled_warmup_lr = ( + linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS + ) linear_scaled_min_lr = linear_scaled_min_lr * config.TRAIN.ACCUMULATION_STEPS config.defrost() config.TRAIN.BASE_LR = linear_scaled_lr @@ -214,7 +276,9 @@ def train_one_epoch(config, model, data_loader, optimizer, epoch, lr_scheduler, config.freeze() os.makedirs(config.OUTPUT, exist_ok=True) - logger = create_logger(output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}") + logger = create_logger( + output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" + ) if dist.get_rank() == 0: path = os.path.join(config.OUTPUT, "config.json") @@ -225,4 +289,4 @@ def train_one_epoch(config, model, data_loader, optimizer, epoch, lr_scheduler, # print config logger.info(config.dump()) - main(config) \ No newline at end of file + main(config) diff --git a/models/__init__.py b/models/__init__.py index 2d9c65e3..59774f75 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -1 +1 @@ -from .build import build_model \ No newline at end of file +from .build import build_model diff --git a/models/build.py b/models/build.py index c37384d2..2d8b07f0 100644 --- a/models/build.py +++ b/models/build.py @@ -5,11 +5,11 @@ # Written by Ze Liu # -------------------------------------------------------- +from .simmim import build_simmim +from .swin_mlp import SwinMLP from .swin_transformer import SwinTransformer -from .swin_transformer_v2 import SwinTransformerV2 from .swin_transformer_moe import SwinTransformerMoE -from .swin_mlp import SwinMLP -from .simmim import build_simmim +from .swin_transformer_v2 import SwinTransformerV2 def build_model(config, is_pretrain=False): @@ -19,102 +19,112 @@ def build_model(config, is_pretrain=False): if config.FUSED_LAYERNORM: try: import apex as amp + layernorm = amp.normalization.FusedLayerNorm except: layernorm = None print("To use FusedLayerNorm, please install apex.") else: import torch.nn as nn + layernorm = nn.LayerNorm if is_pretrain: model = build_simmim(config) return model - if model_type == 'swin': - model = SwinTransformer(img_size=config.DATA.IMG_SIZE, - patch_size=config.MODEL.SWIN.PATCH_SIZE, - in_chans=config.MODEL.SWIN.IN_CHANS, - num_classes=config.MODEL.NUM_CLASSES, - embed_dim=config.MODEL.SWIN.EMBED_DIM, - depths=config.MODEL.SWIN.DEPTHS, - num_heads=config.MODEL.SWIN.NUM_HEADS, - window_size=config.MODEL.SWIN.WINDOW_SIZE, - mlp_ratio=config.MODEL.SWIN.MLP_RATIO, - qkv_bias=config.MODEL.SWIN.QKV_BIAS, - qk_scale=config.MODEL.SWIN.QK_SCALE, - drop_rate=config.MODEL.DROP_RATE, - drop_path_rate=config.MODEL.DROP_PATH_RATE, - ape=config.MODEL.SWIN.APE, - norm_layer=layernorm, - patch_norm=config.MODEL.SWIN.PATCH_NORM, - use_checkpoint=config.TRAIN.USE_CHECKPOINT, - fused_window_process=config.FUSED_WINDOW_PROCESS) - elif model_type == 'swinv2': - model = SwinTransformerV2(img_size=config.DATA.IMG_SIZE, - patch_size=config.MODEL.SWINV2.PATCH_SIZE, - in_chans=config.MODEL.SWINV2.IN_CHANS, - num_classes=config.MODEL.NUM_CLASSES, - embed_dim=config.MODEL.SWINV2.EMBED_DIM, - depths=config.MODEL.SWINV2.DEPTHS, - num_heads=config.MODEL.SWINV2.NUM_HEADS, - window_size=config.MODEL.SWINV2.WINDOW_SIZE, - mlp_ratio=config.MODEL.SWINV2.MLP_RATIO, - qkv_bias=config.MODEL.SWINV2.QKV_BIAS, - drop_rate=config.MODEL.DROP_RATE, - drop_path_rate=config.MODEL.DROP_PATH_RATE, - ape=config.MODEL.SWINV2.APE, - patch_norm=config.MODEL.SWINV2.PATCH_NORM, - use_checkpoint=config.TRAIN.USE_CHECKPOINT, - pretrained_window_sizes=config.MODEL.SWINV2.PRETRAINED_WINDOW_SIZES) - elif model_type == 'swin_moe': - model = SwinTransformerMoE(img_size=config.DATA.IMG_SIZE, - patch_size=config.MODEL.SWIN_MOE.PATCH_SIZE, - in_chans=config.MODEL.SWIN_MOE.IN_CHANS, - num_classes=config.MODEL.NUM_CLASSES, - embed_dim=config.MODEL.SWIN_MOE.EMBED_DIM, - depths=config.MODEL.SWIN_MOE.DEPTHS, - num_heads=config.MODEL.SWIN_MOE.NUM_HEADS, - window_size=config.MODEL.SWIN_MOE.WINDOW_SIZE, - mlp_ratio=config.MODEL.SWIN_MOE.MLP_RATIO, - qkv_bias=config.MODEL.SWIN_MOE.QKV_BIAS, - qk_scale=config.MODEL.SWIN_MOE.QK_SCALE, - drop_rate=config.MODEL.DROP_RATE, - drop_path_rate=config.MODEL.DROP_PATH_RATE, - ape=config.MODEL.SWIN_MOE.APE, - patch_norm=config.MODEL.SWIN_MOE.PATCH_NORM, - mlp_fc2_bias=config.MODEL.SWIN_MOE.MLP_FC2_BIAS, - init_std=config.MODEL.SWIN_MOE.INIT_STD, - use_checkpoint=config.TRAIN.USE_CHECKPOINT, - pretrained_window_sizes=config.MODEL.SWIN_MOE.PRETRAINED_WINDOW_SIZES, - moe_blocks=config.MODEL.SWIN_MOE.MOE_BLOCKS, - num_local_experts=config.MODEL.SWIN_MOE.NUM_LOCAL_EXPERTS, - top_value=config.MODEL.SWIN_MOE.TOP_VALUE, - capacity_factor=config.MODEL.SWIN_MOE.CAPACITY_FACTOR, - cosine_router=config.MODEL.SWIN_MOE.COSINE_ROUTER, - normalize_gate=config.MODEL.SWIN_MOE.NORMALIZE_GATE, - use_bpr=config.MODEL.SWIN_MOE.USE_BPR, - is_gshard_loss=config.MODEL.SWIN_MOE.IS_GSHARD_LOSS, - gate_noise=config.MODEL.SWIN_MOE.GATE_NOISE, - cosine_router_dim=config.MODEL.SWIN_MOE.COSINE_ROUTER_DIM, - cosine_router_init_t=config.MODEL.SWIN_MOE.COSINE_ROUTER_INIT_T, - moe_drop=config.MODEL.SWIN_MOE.MOE_DROP, - aux_loss_weight=config.MODEL.SWIN_MOE.AUX_LOSS_WEIGHT) - elif model_type == 'swin_mlp': - model = SwinMLP(img_size=config.DATA.IMG_SIZE, - patch_size=config.MODEL.SWIN_MLP.PATCH_SIZE, - in_chans=config.MODEL.SWIN_MLP.IN_CHANS, - num_classes=config.MODEL.NUM_CLASSES, - embed_dim=config.MODEL.SWIN_MLP.EMBED_DIM, - depths=config.MODEL.SWIN_MLP.DEPTHS, - num_heads=config.MODEL.SWIN_MLP.NUM_HEADS, - window_size=config.MODEL.SWIN_MLP.WINDOW_SIZE, - mlp_ratio=config.MODEL.SWIN_MLP.MLP_RATIO, - drop_rate=config.MODEL.DROP_RATE, - drop_path_rate=config.MODEL.DROP_PATH_RATE, - ape=config.MODEL.SWIN_MLP.APE, - patch_norm=config.MODEL.SWIN_MLP.PATCH_NORM, - use_checkpoint=config.TRAIN.USE_CHECKPOINT) + if model_type == "swin": + model = SwinTransformer( + img_size=config.DATA.IMG_SIZE, + patch_size=config.MODEL.SWIN.PATCH_SIZE, + in_chans=config.MODEL.SWIN.IN_CHANS, + num_classes=config.MODEL.NUM_CLASSES, + embed_dim=config.MODEL.SWIN.EMBED_DIM, + depths=config.MODEL.SWIN.DEPTHS, + num_heads=config.MODEL.SWIN.NUM_HEADS, + window_size=config.MODEL.SWIN.WINDOW_SIZE, + mlp_ratio=config.MODEL.SWIN.MLP_RATIO, + qkv_bias=config.MODEL.SWIN.QKV_BIAS, + qk_scale=config.MODEL.SWIN.QK_SCALE, + drop_rate=config.MODEL.DROP_RATE, + drop_path_rate=config.MODEL.DROP_PATH_RATE, + ape=config.MODEL.SWIN.APE, + norm_layer=layernorm, + patch_norm=config.MODEL.SWIN.PATCH_NORM, + use_checkpoint=config.TRAIN.USE_CHECKPOINT, + fused_window_process=config.FUSED_WINDOW_PROCESS, + ) + elif model_type == "swinv2": + model = SwinTransformerV2( + img_size=config.DATA.IMG_SIZE, + patch_size=config.MODEL.SWINV2.PATCH_SIZE, + in_chans=config.MODEL.SWINV2.IN_CHANS, + num_classes=config.MODEL.NUM_CLASSES, + embed_dim=config.MODEL.SWINV2.EMBED_DIM, + depths=config.MODEL.SWINV2.DEPTHS, + num_heads=config.MODEL.SWINV2.NUM_HEADS, + window_size=config.MODEL.SWINV2.WINDOW_SIZE, + mlp_ratio=config.MODEL.SWINV2.MLP_RATIO, + qkv_bias=config.MODEL.SWINV2.QKV_BIAS, + drop_rate=config.MODEL.DROP_RATE, + drop_path_rate=config.MODEL.DROP_PATH_RATE, + ape=config.MODEL.SWINV2.APE, + patch_norm=config.MODEL.SWINV2.PATCH_NORM, + use_checkpoint=config.TRAIN.USE_CHECKPOINT, + pretrained_window_sizes=config.MODEL.SWINV2.PRETRAINED_WINDOW_SIZES, + ) + elif model_type == "swin_moe": + model = SwinTransformerMoE( + img_size=config.DATA.IMG_SIZE, + patch_size=config.MODEL.SWIN_MOE.PATCH_SIZE, + in_chans=config.MODEL.SWIN_MOE.IN_CHANS, + num_classes=config.MODEL.NUM_CLASSES, + embed_dim=config.MODEL.SWIN_MOE.EMBED_DIM, + depths=config.MODEL.SWIN_MOE.DEPTHS, + num_heads=config.MODEL.SWIN_MOE.NUM_HEADS, + window_size=config.MODEL.SWIN_MOE.WINDOW_SIZE, + mlp_ratio=config.MODEL.SWIN_MOE.MLP_RATIO, + qkv_bias=config.MODEL.SWIN_MOE.QKV_BIAS, + qk_scale=config.MODEL.SWIN_MOE.QK_SCALE, + drop_rate=config.MODEL.DROP_RATE, + drop_path_rate=config.MODEL.DROP_PATH_RATE, + ape=config.MODEL.SWIN_MOE.APE, + patch_norm=config.MODEL.SWIN_MOE.PATCH_NORM, + mlp_fc2_bias=config.MODEL.SWIN_MOE.MLP_FC2_BIAS, + init_std=config.MODEL.SWIN_MOE.INIT_STD, + use_checkpoint=config.TRAIN.USE_CHECKPOINT, + pretrained_window_sizes=config.MODEL.SWIN_MOE.PRETRAINED_WINDOW_SIZES, + moe_blocks=config.MODEL.SWIN_MOE.MOE_BLOCKS, + num_local_experts=config.MODEL.SWIN_MOE.NUM_LOCAL_EXPERTS, + top_value=config.MODEL.SWIN_MOE.TOP_VALUE, + capacity_factor=config.MODEL.SWIN_MOE.CAPACITY_FACTOR, + cosine_router=config.MODEL.SWIN_MOE.COSINE_ROUTER, + normalize_gate=config.MODEL.SWIN_MOE.NORMALIZE_GATE, + use_bpr=config.MODEL.SWIN_MOE.USE_BPR, + is_gshard_loss=config.MODEL.SWIN_MOE.IS_GSHARD_LOSS, + gate_noise=config.MODEL.SWIN_MOE.GATE_NOISE, + cosine_router_dim=config.MODEL.SWIN_MOE.COSINE_ROUTER_DIM, + cosine_router_init_t=config.MODEL.SWIN_MOE.COSINE_ROUTER_INIT_T, + moe_drop=config.MODEL.SWIN_MOE.MOE_DROP, + aux_loss_weight=config.MODEL.SWIN_MOE.AUX_LOSS_WEIGHT, + ) + elif model_type == "swin_mlp": + model = SwinMLP( + img_size=config.DATA.IMG_SIZE, + patch_size=config.MODEL.SWIN_MLP.PATCH_SIZE, + in_chans=config.MODEL.SWIN_MLP.IN_CHANS, + num_classes=config.MODEL.NUM_CLASSES, + embed_dim=config.MODEL.SWIN_MLP.EMBED_DIM, + depths=config.MODEL.SWIN_MLP.DEPTHS, + num_heads=config.MODEL.SWIN_MLP.NUM_HEADS, + window_size=config.MODEL.SWIN_MLP.WINDOW_SIZE, + mlp_ratio=config.MODEL.SWIN_MLP.MLP_RATIO, + drop_rate=config.MODEL.DROP_RATE, + drop_path_rate=config.MODEL.DROP_PATH_RATE, + ape=config.MODEL.SWIN_MLP.APE, + patch_norm=config.MODEL.SWIN_MLP.PATCH_NORM, + use_checkpoint=config.TRAIN.USE_CHECKPOINT, + ) else: raise NotImplementedError(f"Unkown model: {model_type}") diff --git a/models/simmim.py b/models/simmim.py index fc482b80..731a1c06 100644 --- a/models/simmim.py +++ b/models/simmim.py @@ -1,5 +1,3 @@ - - # -------------------------------------------------------- # SimMIM # Copyright (c) 2021 Microsoft @@ -20,21 +18,41 @@ def norm_targets(targets, patch_size): assert patch_size % 2 == 1 - + targets_ = targets targets_count = torch.ones_like(targets) - targets_square = targets ** 2. - - targets_mean = F.avg_pool2d(targets, kernel_size=patch_size, stride=1, padding=patch_size // 2, count_include_pad=False) - targets_square_mean = F.avg_pool2d(targets_square, kernel_size=patch_size, stride=1, padding=patch_size // 2, count_include_pad=False) - targets_count = F.avg_pool2d(targets_count, kernel_size=patch_size, stride=1, padding=patch_size // 2, count_include_pad=True) * (patch_size ** 2) - - targets_var = (targets_square_mean - targets_mean ** 2.) * (targets_count / (targets_count - 1)) - targets_var = torch.clamp(targets_var, min=0.) - - targets_ = (targets_ - targets_mean) / (targets_var + 1.e-6) ** 0.5 - + targets_square = targets**2.0 + + targets_mean = F.avg_pool2d( + targets, + kernel_size=patch_size, + stride=1, + padding=patch_size // 2, + count_include_pad=False, + ) + targets_square_mean = F.avg_pool2d( + targets_square, + kernel_size=patch_size, + stride=1, + padding=patch_size // 2, + count_include_pad=False, + ) + targets_count = F.avg_pool2d( + targets_count, + kernel_size=patch_size, + stride=1, + padding=patch_size // 2, + count_include_pad=True, + ) * (patch_size**2) + + targets_var = (targets_square_mean - targets_mean**2.0) * ( + targets_count / (targets_count - 1) + ) + targets_var = torch.clamp(targets_var, min=0.0) + + targets_ = (targets_ - targets_mean) / (targets_var + 1.0e-6) ** 0.5 + return targets_ @@ -45,7 +63,7 @@ def __init__(self, **kwargs): assert self.num_classes == 0 self.mask_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim)) - trunc_normal_(self.mask_token, mean=0., std=.02) + trunc_normal_(self.mask_token, mean=0.0, std=0.02) def forward(self, x, mask): x = self.patch_embed(x) @@ -55,7 +73,7 @@ def forward(self, x, mask): mask_tokens = self.mask_token.expand(B, L, -1) w = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) - x = x * (1. - w) + mask_tokens * w + x = x * (1.0 - w) + mask_tokens * w if self.ape: x = x + self.absolute_pos_embed @@ -67,13 +85,13 @@ def forward(self, x, mask): x = x.transpose(1, 2) B, C, L = x.shape - H = W = int(L ** 0.5) + H = W = int(L**0.5) x = x.reshape(B, C, H, W) return x @torch.jit.ignore def no_weight_decay(self): - return super().no_weight_decay() | {'mask_token'} + return super().no_weight_decay() | {"mask_token"} class SwinTransformerV2ForSimMIM(SwinTransformerV2): @@ -83,7 +101,7 @@ def __init__(self, **kwargs): assert self.num_classes == 0 self.mask_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim)) - trunc_normal_(self.mask_token, mean=0., std=.02) + trunc_normal_(self.mask_token, mean=0.0, std=0.02) def forward(self, x, mask): x = self.patch_embed(x) @@ -93,7 +111,7 @@ def forward(self, x, mask): mask_tokens = self.mask_token.expand(B, L, -1) w = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) - x = x * (1. - w) + mask_tokens * w + x = x * (1.0 - w) + mask_tokens * w if self.ape: x = x + self.absolute_pos_embed @@ -105,13 +123,13 @@ def forward(self, x, mask): x = x.transpose(1, 2) B, C, L = x.shape - H = W = int(L ** 0.5) + H = W = int(L**0.5) x = x.reshape(B, C, H, W) return x @torch.jit.ignore def no_weight_decay(self): - return super().no_weight_decay() | {'mask_token'} + return super().no_weight_decay() | {"mask_token"} class SimMIM(nn.Module): @@ -124,7 +142,9 @@ def __init__(self, config, encoder, encoder_stride, in_chans, patch_size): self.decoder = nn.Sequential( nn.Conv2d( in_channels=self.encoder.num_features, - out_channels=self.encoder_stride ** 2 * 3, kernel_size=1), + out_channels=self.encoder_stride**2 * 3, + kernel_size=1, + ), nn.PixelShuffle(self.encoder_stride), ) @@ -135,32 +155,37 @@ def forward(self, x, mask): z = self.encoder(x, mask) x_rec = self.decoder(z) - mask = mask.repeat_interleave(self.patch_size, 1).repeat_interleave(self.patch_size, 2).unsqueeze(1).contiguous() - + mask = ( + mask.repeat_interleave(self.patch_size, 1) + .repeat_interleave(self.patch_size, 2) + .unsqueeze(1) + .contiguous() + ) + # norm target as prompted if self.config.NORM_TARGET.ENABLE: x = norm_targets(x, self.config.NORM_TARGET.PATCH_SIZE) - - loss_recon = F.l1_loss(x, x_rec, reduction='none') + + loss_recon = F.l1_loss(x, x_rec, reduction="none") loss = (loss_recon * mask).sum() / (mask.sum() + 1e-5) / self.in_chans return loss @torch.jit.ignore def no_weight_decay(self): - if hasattr(self.encoder, 'no_weight_decay'): - return {'encoder.' + i for i in self.encoder.no_weight_decay()} + if hasattr(self.encoder, "no_weight_decay"): + return {"encoder." + i for i in self.encoder.no_weight_decay()} return {} @torch.jit.ignore def no_weight_decay_keywords(self): - if hasattr(self.encoder, 'no_weight_decay_keywords'): - return {'encoder.' + i for i in self.encoder.no_weight_decay_keywords()} + if hasattr(self.encoder, "no_weight_decay_keywords"): + return {"encoder." + i for i in self.encoder.no_weight_decay_keywords()} return {} def build_simmim(config): model_type = config.MODEL.TYPE - if model_type == 'swin': + if model_type == "swin": encoder = SwinTransformerForSimMIM( img_size=config.DATA.IMG_SIZE, patch_size=config.MODEL.SWIN.PATCH_SIZE, @@ -177,11 +202,12 @@ def build_simmim(config): drop_path_rate=config.MODEL.DROP_PATH_RATE, ape=config.MODEL.SWIN.APE, patch_norm=config.MODEL.SWIN.PATCH_NORM, - use_checkpoint=config.TRAIN.USE_CHECKPOINT) + use_checkpoint=config.TRAIN.USE_CHECKPOINT, + ) encoder_stride = 32 in_chans = config.MODEL.SWIN.IN_CHANS patch_size = config.MODEL.SWIN.PATCH_SIZE - elif model_type == 'swinv2': + elif model_type == "swinv2": encoder = SwinTransformerV2ForSimMIM( img_size=config.DATA.IMG_SIZE, patch_size=config.MODEL.SWINV2.PATCH_SIZE, @@ -197,13 +223,20 @@ def build_simmim(config): drop_path_rate=config.MODEL.DROP_PATH_RATE, ape=config.MODEL.SWINV2.APE, patch_norm=config.MODEL.SWINV2.PATCH_NORM, - use_checkpoint=config.TRAIN.USE_CHECKPOINT) + use_checkpoint=config.TRAIN.USE_CHECKPOINT, + ) encoder_stride = 32 in_chans = config.MODEL.SWINV2.IN_CHANS patch_size = config.MODEL.SWINV2.PATCH_SIZE else: raise NotImplementedError(f"Unknown pre-train model: {model_type}") - model = SimMIM(config=config.MODEL.SIMMIM, encoder=encoder, encoder_stride=encoder_stride, in_chans=in_chans, patch_size=patch_size) + model = SimMIM( + config=config.MODEL.SIMMIM, + encoder=encoder, + encoder_stride=encoder_stride, + in_chans=in_chans, + patch_size=patch_size, + ) - return model \ No newline at end of file + return model diff --git a/models/swin_mlp.py b/models/swin_mlp.py index 115c43cd..bd0ae21b 100644 --- a/models/swin_mlp.py +++ b/models/swin_mlp.py @@ -13,7 +13,14 @@ class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + drop=0.0, + ): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -42,7 +49,9 @@ def window_partition(x, window_size): """ B, H, W, C = x.shape x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) - windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + windows = ( + x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + ) return windows @@ -58,13 +67,15 @@ def window_reverse(windows, window_size, H, W): x: (B, H, W, C) """ B = int(windows.shape[0] / (H * W / window_size / window_size)) - x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = windows.view( + B, H // window_size, W // window_size, window_size, window_size, -1 + ) x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) return x class SwinMLPBlock(nn.Module): - r""" Swin MLP Block. + r"""Swin MLP Block. Args: dim (int): Number of input channels. @@ -79,9 +90,19 @@ class SwinMLPBlock(nn.Module): norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm """ - def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, - mlp_ratio=4., drop=0., drop_path=0., - act_layer=nn.GELU, norm_layer=nn.LayerNorm): + def __init__( + self, + dim, + input_resolution, + num_heads, + window_size=7, + shift_size=0, + mlp_ratio=4.0, + drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + ): super().__init__() self.dim = dim self.input_resolution = input_resolution @@ -93,22 +114,35 @@ def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0 # if window size is larger than input resolution, we don't partition windows self.shift_size = 0 self.window_size = min(self.input_resolution) - assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + assert ( + 0 <= self.shift_size < self.window_size + ), "shift_size must in 0-window_size" - self.padding = [self.window_size - self.shift_size, self.shift_size, - self.window_size - self.shift_size, self.shift_size] # P_l,P_r,P_t,P_b + self.padding = [ + self.window_size - self.shift_size, + self.shift_size, + self.window_size - self.shift_size, + self.shift_size, + ] # P_l,P_r,P_t,P_b self.norm1 = norm_layer(dim) # use group convolution to implement multi-head MLP - self.spatial_mlp = nn.Conv1d(self.num_heads * self.window_size ** 2, - self.num_heads * self.window_size ** 2, - kernel_size=1, - groups=self.num_heads) - - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.spatial_mlp = nn.Conv1d( + self.num_heads * self.window_size**2, + self.num_heads * self.window_size**2, + kernel_size=1, + groups=self.num_heads, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + ) def forward(self, x): H, W = self.input_resolution @@ -128,22 +162,42 @@ def forward(self, x): _, _H, _W, _ = shifted_x.shape # partition windows - x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C - x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + x_windows = window_partition( + shifted_x, self.window_size + ) # nW*B, window_size, window_size, C + x_windows = x_windows.view( + -1, self.window_size * self.window_size, C + ) # nW*B, window_size*window_size, C # Window/Shifted-Window Spatial MLP - x_windows_heads = x_windows.view(-1, self.window_size * self.window_size, self.num_heads, C // self.num_heads) - x_windows_heads = x_windows_heads.transpose(1, 2) # nW*B, nH, window_size*window_size, C//nH - x_windows_heads = x_windows_heads.reshape(-1, self.num_heads * self.window_size * self.window_size, - C // self.num_heads) - spatial_mlp_windows = self.spatial_mlp(x_windows_heads) # nW*B, nH*window_size*window_size, C//nH - spatial_mlp_windows = spatial_mlp_windows.view(-1, self.num_heads, self.window_size * self.window_size, - C // self.num_heads).transpose(1, 2) - spatial_mlp_windows = spatial_mlp_windows.reshape(-1, self.window_size * self.window_size, C) + x_windows_heads = x_windows.view( + -1, self.window_size * self.window_size, self.num_heads, C // self.num_heads + ) + x_windows_heads = x_windows_heads.transpose( + 1, 2 + ) # nW*B, nH, window_size*window_size, C//nH + x_windows_heads = x_windows_heads.reshape( + -1, + self.num_heads * self.window_size * self.window_size, + C // self.num_heads, + ) + spatial_mlp_windows = self.spatial_mlp( + x_windows_heads + ) # nW*B, nH*window_size*window_size, C//nH + spatial_mlp_windows = spatial_mlp_windows.view( + -1, self.num_heads, self.window_size * self.window_size, C // self.num_heads + ).transpose(1, 2) + spatial_mlp_windows = spatial_mlp_windows.reshape( + -1, self.window_size * self.window_size, C + ) # merge windows - spatial_mlp_windows = spatial_mlp_windows.reshape(-1, self.window_size, self.window_size, C) - shifted_x = window_reverse(spatial_mlp_windows, self.window_size, _H, _W) # B H' W' C + spatial_mlp_windows = spatial_mlp_windows.reshape( + -1, self.window_size, self.window_size, C + ) + shifted_x = window_reverse( + spatial_mlp_windows, self.window_size, _H, _W + ) # B H' W' C # reverse shift if self.shift_size > 0: @@ -160,8 +214,10 @@ def forward(self, x): return x def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ - f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + return ( + f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " + f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + ) def flops(self): flops = 0 @@ -174,7 +230,12 @@ def flops(self): nW = (H / self.window_size + 1) * (W / self.window_size + 1) else: nW = H * W / self.window_size / self.window_size - flops += nW * self.dim * (self.window_size * self.window_size) * (self.window_size * self.window_size) + flops += ( + nW + * self.dim + * (self.window_size * self.window_size) + * (self.window_size * self.window_size) + ) # mlp flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio # norm2 @@ -183,7 +244,7 @@ def flops(self): class PatchMerging(nn.Module): - r""" Patch Merging Layer. + r"""Patch Merging Layer. Args: input_resolution (tuple[int]): Resolution of input feature. @@ -232,7 +293,7 @@ def flops(self): class BasicLayer(nn.Module): - """ A basic Swin MLP layer for one stage. + """A basic Swin MLP layer for one stage. Args: dim (int): Number of input channels. @@ -248,9 +309,20 @@ class BasicLayer(nn.Module): use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. """ - def __init__(self, dim, input_resolution, depth, num_heads, window_size, - mlp_ratio=4., drop=0., drop_path=0., - norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False): + def __init__( + self, + dim, + input_resolution, + depth, + num_heads, + window_size, + mlp_ratio=4.0, + drop=0.0, + drop_path=0.0, + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False, + ): super().__init__() self.dim = dim @@ -259,19 +331,30 @@ def __init__(self, dim, input_resolution, depth, num_heads, window_size, self.use_checkpoint = use_checkpoint # build blocks - self.blocks = nn.ModuleList([ - SwinMLPBlock(dim=dim, input_resolution=input_resolution, - num_heads=num_heads, window_size=window_size, - shift_size=0 if (i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - drop=drop, - drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, - norm_layer=norm_layer) - for i in range(depth)]) + self.blocks = nn.ModuleList( + [ + SwinMLPBlock( + dim=dim, + input_resolution=input_resolution, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + drop=drop, + drop_path=drop_path[i] + if isinstance(drop_path, list) + else drop_path, + norm_layer=norm_layer, + ) + for i in range(depth) + ] + ) # patch merging layer if downsample is not None: - self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer) + self.downsample = downsample( + input_resolution, dim=dim, norm_layer=norm_layer + ) else: self.downsample = None @@ -298,7 +381,7 @@ def flops(self): class PatchEmbed(nn.Module): - r""" Image to Patch Embedding + r"""Image to Patch Embedding Args: img_size (int): Image size. Default: 224. @@ -308,11 +391,16 @@ class PatchEmbed(nn.Module): norm_layer (nn.Module, optional): Normalization layer. Default: None """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + def __init__( + self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None + ): super().__init__() img_size = to_2tuple(img_size) patch_size = to_2tuple(patch_size) - patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] + patches_resolution = [ + img_size[0] // patch_size[0], + img_size[1] // patch_size[1], + ] self.img_size = img_size self.patch_size = patch_size self.patches_resolution = patches_resolution @@ -321,7 +409,9 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_la self.in_chans = in_chans self.embed_dim = embed_dim - self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=patch_size, stride=patch_size + ) if norm_layer is not None: self.norm = norm_layer(embed_dim) else: @@ -330,8 +420,9 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_la def forward(self, x): B, C, H, W = x.shape # FIXME look at relaxing size constraints - assert H == self.img_size[0] and W == self.img_size[1], \ - f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." + assert ( + H == self.img_size[0] and W == self.img_size[1] + ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C if self.norm is not None: x = self.norm(x) @@ -339,14 +430,20 @@ def forward(self, x): def flops(self): Ho, Wo = self.patches_resolution - flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + flops = ( + Ho + * Wo + * self.embed_dim + * self.in_chans + * (self.patch_size[0] * self.patch_size[1]) + ) if self.norm is not None: flops += Ho * Wo * self.embed_dim return flops class SwinMLP(nn.Module): - r""" Swin MLP + r"""Swin MLP Args: img_size (int | tuple(int)): Input image size. Default 224 @@ -366,11 +463,25 @@ class SwinMLP(nn.Module): use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, - embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], - window_size=7, mlp_ratio=4., drop_rate=0., drop_path_rate=0.1, - norm_layer=nn.LayerNorm, ape=False, patch_norm=True, - use_checkpoint=False, **kwargs): + def __init__( + self, + img_size=224, + patch_size=4, + in_chans=3, + num_classes=1000, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4.0, + drop_rate=0.0, + drop_path_rate=0.1, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + use_checkpoint=False, + **kwargs, + ): super().__init__() self.num_classes = num_classes @@ -383,48 +494,64 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, # split image into non-overlapping patches self.patch_embed = PatchEmbed( - img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None) + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, + ) num_patches = self.patch_embed.num_patches patches_resolution = self.patch_embed.patches_resolution self.patches_resolution = patches_resolution # absolute position embedding if self.ape: - self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) - trunc_normal_(self.absolute_pos_embed, std=.02) + self.absolute_pos_embed = nn.Parameter( + torch.zeros(1, num_patches, embed_dim) + ) + trunc_normal_(self.absolute_pos_embed, std=0.02) self.pos_drop = nn.Dropout(p=drop_rate) # stochastic depth - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule # build layers self.layers = nn.ModuleList() for i_layer in range(self.num_layers): - layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer), - input_resolution=(patches_resolution[0] // (2 ** i_layer), - patches_resolution[1] // (2 ** i_layer)), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - drop=drop_rate, - drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], - norm_layer=norm_layer, - downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, - use_checkpoint=use_checkpoint) + layer = BasicLayer( + dim=int(embed_dim * 2**i_layer), + input_resolution=( + patches_resolution[0] // (2**i_layer), + patches_resolution[1] // (2**i_layer), + ), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=self.mlp_ratio, + drop=drop_rate, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint, + ) self.layers.append(layer) self.norm = norm_layer(self.num_features) self.avgpool = nn.AdaptiveAvgPool1d(1) - self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity() + self.head = ( + nn.Linear(self.num_features, num_classes) + if num_classes > 0 + else nn.Identity() + ) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, (nn.Linear, nn.Conv1d)): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -433,11 +560,11 @@ def _init_weights(self, m): @torch.jit.ignore def no_weight_decay(self): - return {'absolute_pos_embed'} + return {"absolute_pos_embed"} @torch.jit.ignore def no_weight_decay_keywords(self): - return {'relative_position_bias_table'} + return {"relative_position_bias_table"} def forward_features(self, x): x = self.patch_embed(x) @@ -463,6 +590,11 @@ def flops(self): flops += self.patch_embed.flops() for i, layer in enumerate(self.layers): flops += layer.flops() - flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers) + flops += ( + self.num_features + * self.patches_resolution[0] + * self.patches_resolution[1] + // (2**self.num_layers) + ) flops += self.num_features * self.num_classes return flops diff --git a/models/swin_transformer.py b/models/swin_transformer.py index dde06bc5..a7ba5f15 100644 --- a/models/swin_transformer.py +++ b/models/swin_transformer.py @@ -11,20 +11,33 @@ from timm.models.layers import DropPath, to_2tuple, trunc_normal_ try: - import os, sys + import os + import sys - kernel_path = os.path.abspath(os.path.join('..')) + kernel_path = os.path.abspath(os.path.join("..")) sys.path.append(kernel_path) - from kernels.window_process.window_process import WindowProcess, WindowProcessReverse + from kernels.window_process.window_process import ( + WindowProcess, + WindowProcessReverse, + ) except: WindowProcess = None WindowProcessReverse = None - print("[Warning] Fused window process have not been installed. Please refer to get_started.md for installation.") + print( + "[Warning] Fused window process have not been installed. Please refer to get_started.md for installation." + ) class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + drop=0.0, + ): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -53,7 +66,9 @@ def window_partition(x, window_size): """ B, H, W, C = x.shape x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) - windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + windows = ( + x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + ) return windows @@ -69,13 +84,15 @@ def window_reverse(windows, window_size, H, W): x: (B, H, W, C) """ B = int(windows.shape[0] / (H * W / window_size / window_size)) - x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = windows.view( + B, H // window_size, W // window_size, window_size, window_size, -1 + ) x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) return x class WindowAttention(nn.Module): - r""" Window based multi-head self attention (W-MSA) module with relative position bias. + r"""Window based multi-head self attention (W-MSA) module with relative position bias. It supports both of shifted and non-shifted window. Args: @@ -88,26 +105,40 @@ class WindowAttention(nn.Module): proj_drop (float, optional): Dropout ratio of output. Default: 0.0 """ - def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + def __init__( + self, + dim, + window_size, + num_heads, + qkv_bias=True, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + ): super().__init__() self.dim = dim self.window_size = window_size # Wh, Ww self.num_heads = num_heads head_dim = dim // num_heads - self.scale = qk_scale or head_dim ** -0.5 + self.scale = qk_scale or head_dim**-0.5 # define a parameter table of relative position bias self.relative_position_bias_table = nn.Parameter( - torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) + ) # 2*Wh-1 * 2*Ww-1, nH # get pair-wise relative position index for each token inside the window coords_h = torch.arange(self.window_size[0]) coords_w = torch.arange(self.window_size[1]) coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww - relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww - relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords = ( + coords_flatten[:, :, None] - coords_flatten[:, None, :] + ) # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute( + 1, 2, 0 + ).contiguous() # Wh*Ww, Wh*Ww, 2 relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 relative_coords[:, :, 1] += self.window_size[1] - 1 relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 @@ -119,7 +150,7 @@ def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, at self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) - trunc_normal_(self.relative_position_bias_table, std=.02) + trunc_normal_(self.relative_position_bias_table, std=0.02) self.softmax = nn.Softmax(dim=-1) def forward(self, x, mask=None): @@ -129,20 +160,37 @@ def forward(self, x, mask=None): mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None """ B_, N, C = x.shape - qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) - q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + qkv = ( + self.qkv(x) + .reshape(B_, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + q, k, v = ( + qkv[0], + qkv[1], + qkv[2], + ) # make torchscript happy (cannot use tensor as tuple) q = q * self.scale - attn = (q @ k.transpose(-2, -1)) - - relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( - self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH - relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = q @ k.transpose(-2, -1) + + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1], + self.window_size[0] * self.window_size[1], + -1, + ) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() # nH, Wh*Ww, Wh*Ww attn = attn + relative_position_bias.unsqueeze(0) if mask is not None: nW = mask.shape[0] - attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze( + 1 + ).unsqueeze(0) attn = attn.view(-1, self.num_heads, N, N) attn = self.softmax(attn) else: @@ -156,7 +204,7 @@ def forward(self, x, mask=None): return x def extra_repr(self) -> str: - return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}' + return f"dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}" def flops(self, N): # calculate flops for 1 window with token length of N @@ -173,7 +221,7 @@ def flops(self, N): class SwinTransformerBlock(nn.Module): - r""" Swin Transformer Block. + r"""Swin Transformer Block. Args: dim (int): Number of input channels. @@ -192,10 +240,23 @@ class SwinTransformerBlock(nn.Module): fused_window_process (bool, optional): If True, use one kernel to fused window shift & window partition for acceleration, similar for the reversed part. Default: False """ - def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., - act_layer=nn.GELU, norm_layer=nn.LayerNorm, - fused_window_process=False): + def __init__( + self, + dim, + input_resolution, + num_heads, + window_size=7, + shift_size=0, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + fused_window_process=False, + ): super().__init__() self.dim = dim self.input_resolution = input_resolution @@ -207,38 +268,59 @@ def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0 # if window size is larger than input resolution, we don't partition windows self.shift_size = 0 self.window_size = min(self.input_resolution) - assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + assert ( + 0 <= self.shift_size < self.window_size + ), "shift_size must in 0-window_size" self.norm1 = norm_layer(dim) self.attn = WindowAttention( - dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, - qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + dim, + window_size=to_2tuple(self.window_size), + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + ) if self.shift_size > 0: # calculate attention mask for SW-MSA H, W = self.input_resolution img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 - h_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) - w_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) + h_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + w_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 - mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = window_partition( + img_mask, self.window_size + ) # nW, window_size, window_size, 1 mask_windows = mask_windows.view(-1, self.window_size * self.window_size) attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + attn_mask = attn_mask.masked_fill( + attn_mask != 0, float(-100.0) + ).masked_fill(attn_mask == 0, float(0.0)) else: attn_mask = None @@ -257,20 +339,32 @@ def forward(self, x): # cyclic shift if self.shift_size > 0: if not self.fused_window_process: - shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + shifted_x = torch.roll( + x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2) + ) # partition windows - x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = window_partition( + shifted_x, self.window_size + ) # nW*B, window_size, window_size, C else: - x_windows = WindowProcess.apply(x, B, H, W, C, -self.shift_size, self.window_size) + x_windows = WindowProcess.apply( + x, B, H, W, C, -self.shift_size, self.window_size + ) else: shifted_x = x # partition windows - x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = window_partition( + shifted_x, self.window_size + ) # nW*B, window_size, window_size, C - x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + x_windows = x_windows.view( + -1, self.window_size * self.window_size, C + ) # nW*B, window_size*window_size, C # W-MSA/SW-MSA - attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C + attn_windows = self.attn( + x_windows, mask=self.attn_mask + ) # nW*B, window_size*window_size, C # merge windows attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) @@ -278,12 +372,20 @@ def forward(self, x): # reverse cyclic shift if self.shift_size > 0: if not self.fused_window_process: - shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C - x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + shifted_x = window_reverse( + attn_windows, self.window_size, H, W + ) # B H' W' C + x = torch.roll( + shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2) + ) else: - x = WindowProcessReverse.apply(attn_windows, B, H, W, C, self.shift_size, self.window_size) + x = WindowProcessReverse.apply( + attn_windows, B, H, W, C, self.shift_size, self.window_size + ) else: - shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C + shifted_x = window_reverse( + attn_windows, self.window_size, H, W + ) # B H' W' C x = shifted_x x = x.view(B, H * W, C) x = shortcut + self.drop_path(x) @@ -294,8 +396,10 @@ def forward(self, x): return x def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ - f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + return ( + f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " + f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + ) def flops(self): flops = 0 @@ -313,7 +417,7 @@ def flops(self): class PatchMerging(nn.Module): - r""" Patch Merging Layer. + r"""Patch Merging Layer. Args: input_resolution (tuple[int]): Resolution of input feature. @@ -362,7 +466,7 @@ def flops(self): class BasicLayer(nn.Module): - """ A basic Swin Transformer layer for one stage. + """A basic Swin Transformer layer for one stage. Args: dim (int): Number of input channels. @@ -382,10 +486,24 @@ class BasicLayer(nn.Module): fused_window_process (bool, optional): If True, use one kernel to fused window shift & window partition for acceleration, similar for the reversed part. Default: False """ - def __init__(self, dim, input_resolution, depth, num_heads, window_size, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False, - fused_window_process=False): + def __init__( + self, + dim, + input_resolution, + depth, + num_heads, + window_size, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False, + fused_window_process=False, + ): super().__init__() self.dim = dim @@ -394,21 +512,34 @@ def __init__(self, dim, input_resolution, depth, num_heads, window_size, self.use_checkpoint = use_checkpoint # build blocks - self.blocks = nn.ModuleList([ - SwinTransformerBlock(dim=dim, input_resolution=input_resolution, - num_heads=num_heads, window_size=window_size, - shift_size=0 if (i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop, attn_drop=attn_drop, - drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, - norm_layer=norm_layer, - fused_window_process=fused_window_process) - for i in range(depth)]) + self.blocks = nn.ModuleList( + [ + SwinTransformerBlock( + dim=dim, + input_resolution=input_resolution, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] + if isinstance(drop_path, list) + else drop_path, + norm_layer=norm_layer, + fused_window_process=fused_window_process, + ) + for i in range(depth) + ] + ) # patch merging layer if downsample is not None: - self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer) + self.downsample = downsample( + input_resolution, dim=dim, norm_layer=norm_layer + ) else: self.downsample = None @@ -435,7 +566,7 @@ def flops(self): class PatchEmbed(nn.Module): - r""" Image to Patch Embedding + r"""Image to Patch Embedding Args: img_size (int): Image size. Default: 224. @@ -445,11 +576,16 @@ class PatchEmbed(nn.Module): norm_layer (nn.Module, optional): Normalization layer. Default: None """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + def __init__( + self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None + ): super().__init__() img_size = to_2tuple(img_size) patch_size = to_2tuple(patch_size) - patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] + patches_resolution = [ + img_size[0] // patch_size[0], + img_size[1] // patch_size[1], + ] self.img_size = img_size self.patch_size = patch_size self.patches_resolution = patches_resolution @@ -458,7 +594,9 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_la self.in_chans = in_chans self.embed_dim = embed_dim - self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=patch_size, stride=patch_size + ) if norm_layer is not None: self.norm = norm_layer(embed_dim) else: @@ -467,8 +605,9 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_la def forward(self, x): B, C, H, W = x.shape # FIXME look at relaxing size constraints - assert H == self.img_size[0] and W == self.img_size[1], \ - f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." + assert ( + H == self.img_size[0] and W == self.img_size[1] + ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C if self.norm is not None: x = self.norm(x) @@ -476,14 +615,20 @@ def forward(self, x): def flops(self): Ho, Wo = self.patches_resolution - flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + flops = ( + Ho + * Wo + * self.embed_dim + * self.in_chans + * (self.patch_size[0] * self.patch_size[1]) + ) if self.norm is not None: flops += Ho * Wo * self.embed_dim return flops class SwinTransformer(nn.Module): - r""" Swin Transformer + r"""Swin Transformer A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - https://arxiv.org/pdf/2103.14030 @@ -509,12 +654,29 @@ class SwinTransformer(nn.Module): fused_window_process (bool, optional): If True, use one kernel to fused window shift & window partition for acceleration, similar for the reversed part. Default: False """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, - embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], - window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None, - drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1, - norm_layer=nn.LayerNorm, ape=False, patch_norm=True, - use_checkpoint=False, fused_window_process=False, **kwargs): + def __init__( + self, + img_size=224, + patch_size=4, + in_chans=3, + num_classes=1000, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.1, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + use_checkpoint=False, + fused_window_process=False, + **kwargs, + ): super().__init__() self.num_classes = num_classes @@ -527,50 +689,68 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, # split image into non-overlapping patches self.patch_embed = PatchEmbed( - img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None) + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, + ) num_patches = self.patch_embed.num_patches patches_resolution = self.patch_embed.patches_resolution self.patches_resolution = patches_resolution # absolute position embedding if self.ape: - self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) - trunc_normal_(self.absolute_pos_embed, std=.02) + self.absolute_pos_embed = nn.Parameter( + torch.zeros(1, num_patches, embed_dim) + ) + trunc_normal_(self.absolute_pos_embed, std=0.02) self.pos_drop = nn.Dropout(p=drop_rate) # stochastic depth - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule # build layers self.layers = nn.ModuleList() for i_layer in range(self.num_layers): - layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer), - input_resolution=(patches_resolution[0] // (2 ** i_layer), - patches_resolution[1] // (2 ** i_layer)), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, - drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], - norm_layer=norm_layer, - downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, - use_checkpoint=use_checkpoint, - fused_window_process=fused_window_process) + layer = BasicLayer( + dim=int(embed_dim * 2**i_layer), + input_resolution=( + patches_resolution[0] // (2**i_layer), + patches_resolution[1] // (2**i_layer), + ), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=self.mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint, + fused_window_process=fused_window_process, + ) self.layers.append(layer) self.norm = norm_layer(self.num_features) self.avgpool = nn.AdaptiveAvgPool1d(1) - self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity() + self.head = ( + nn.Linear(self.num_features, num_classes) + if num_classes > 0 + else nn.Identity() + ) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -579,11 +759,11 @@ def _init_weights(self, m): @torch.jit.ignore def no_weight_decay(self): - return {'absolute_pos_embed'} + return {"absolute_pos_embed"} @torch.jit.ignore def no_weight_decay_keywords(self): - return {'relative_position_bias_table'} + return {"relative_position_bias_table"} def forward_features(self, x): x = self.patch_embed(x) @@ -609,6 +789,11 @@ def flops(self): flops += self.patch_embed.flops() for i, layer in enumerate(self.layers): flops += layer.flops() - flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers) + flops += ( + self.num_features + * self.patches_resolution[0] + * self.patches_resolution[1] + // (2**self.num_layers) + ) flops += self.num_features * self.num_classes return flops diff --git a/models/swin_transformer_moe.py b/models/swin_transformer_moe.py index e9f26d43..a714a029 100644 --- a/models/swin_transformer_moe.py +++ b/models/swin_transformer_moe.py @@ -5,24 +5,33 @@ # Written by Ze Liu # -------------------------------------------------------- +import numpy as np import torch +import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F -import torch.distributed as dist import torch.utils.checkpoint as checkpoint from timm.models.layers import DropPath, to_2tuple, trunc_normal_ -import numpy as np try: from tutel import moe as tutel_moe except: tutel_moe = None - print("Tutel has not been installed. To use Swin-MoE, please install Tutel; otherwise, just ignore this.") + print( + "Tutel has not been installed. To use Swin-MoE, please install Tutel; otherwise, just ignore this." + ) class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0., - mlp_fc2_bias=True): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + drop=0.0, + mlp_fc2_bias=True, + ): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -41,10 +50,24 @@ def forward(self, x): class MoEMlp(nn.Module): - def __init__(self, in_features, hidden_features, num_local_experts, top_value, capacity_factor=1.25, - cosine_router=False, normalize_gate=False, use_bpr=True, is_gshard_loss=True, - gate_noise=1.0, cosine_router_dim=256, cosine_router_init_t=0.5, moe_drop=0.0, init_std=0.02, - mlp_fc2_bias=True): + def __init__( + self, + in_features, + hidden_features, + num_local_experts, + top_value, + capacity_factor=1.25, + cosine_router=False, + normalize_gate=False, + use_bpr=True, + is_gshard_loss=True, + gate_noise=1.0, + cosine_router_dim=256, + cosine_router_init_t=0.5, + moe_drop=0.0, + init_std=0.02, + mlp_fc2_bias=True, + ): super().__init__() self.in_features = in_features @@ -62,23 +85,30 @@ def __init__(self, in_features, hidden_features, num_local_experts, top_value, c self._dropout = nn.Dropout(p=moe_drop) - _gate_type = {'type': 'cosine_top' if cosine_router else 'top', - 'k': top_value, 'capacity_factor': capacity_factor, - 'gate_noise': gate_noise, 'fp32_gate': True} + _gate_type = { + "type": "cosine_top" if cosine_router else "top", + "k": top_value, + "capacity_factor": capacity_factor, + "gate_noise": gate_noise, + "fp32_gate": True, + } if cosine_router: - _gate_type['proj_dim'] = cosine_router_dim - _gate_type['init_t'] = cosine_router_init_t + _gate_type["proj_dim"] = cosine_router_dim + _gate_type["init_t"] = cosine_router_init_t self._moe_layer = tutel_moe.moe_layer( gate_type=_gate_type, model_dim=in_features, - experts={'type': 'ffn', 'count_per_node': num_local_experts, 'hidden_size_per_expert': hidden_features, - 'activation_fn': lambda x: self._dropout(F.gelu(x))}, - scan_expert_func=lambda name, param: setattr(param, 'skip_allreduce', True), + experts={ + "type": "ffn", + "count_per_node": num_local_experts, + "hidden_size_per_expert": hidden_features, + "activation_fn": lambda x: self._dropout(F.gelu(x)), + }, + scan_expert_func=lambda name, param: setattr(param, "skip_allreduce", True), seeds=(1, self.dist_rank + 1, self.dist_rank + 1), batch_prioritized_routing=use_bpr, normalize_gate=normalize_gate, is_gshard_loss=is_gshard_loss, - ) if not self.mlp_fc2_bias: self._moe_layer.experts.batched_fc2_bias.requires_grad = False @@ -88,10 +118,12 @@ def forward(self, x): return x, x.l_aux def extra_repr(self) -> str: - return f'[Statistics-{self.dist_rank}] param count for MoE, ' \ - f'in_features = {self.in_features}, hidden_features = {self.hidden_features}, ' \ - f'num_local_experts = {self.num_local_experts}, top_value = {self.top_value}, ' \ - f'cosine_router={self.cosine_router} normalize_gate={self.normalize_gate}, use_bpr = {self.use_bpr}' + return ( + f"[Statistics-{self.dist_rank}] param count for MoE, " + f"in_features = {self.in_features}, hidden_features = {self.hidden_features}, " + f"num_local_experts = {self.num_local_experts}, top_value = {self.top_value}, " + f"cosine_router={self.cosine_router} normalize_gate={self.normalize_gate}, use_bpr = {self.use_bpr}" + ) def _init_weights(self): if hasattr(self._moe_layer, "experts"): @@ -112,7 +144,9 @@ def window_partition(x, window_size): """ B, H, W, C = x.shape x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) - windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + windows = ( + x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + ) return windows @@ -128,13 +162,15 @@ def window_reverse(windows, window_size, H, W): x: (B, H, W, C) """ B = int(windows.shape[0] / (H * W / window_size / window_size)) - x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = windows.view( + B, H // window_size, W // window_size, window_size, window_size, -1 + ) x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) return x class WindowAttention(nn.Module): - r""" Window based multi-head self attention (W-MSA) module with relative position bias. + r"""Window based multi-head self attention (W-MSA) module with relative position bias. It supports both of shifted and non-shifted window. Args: @@ -148,8 +184,17 @@ class WindowAttention(nn.Module): pretrained_window_size (tuple[int]): The height and width of the window in pre-training. """ - def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0., - pretrained_window_size=[0, 0]): + def __init__( + self, + dim, + window_size, + num_heads, + qkv_bias=True, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + pretrained_window_size=[0, 0], + ): super().__init__() self.dim = dim @@ -158,28 +203,40 @@ def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, at self.num_heads = num_heads head_dim = dim // num_heads - self.scale = qk_scale or head_dim ** -0.5 + self.scale = qk_scale or head_dim**-0.5 # mlp to generate continuous relative position bias - self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True), - nn.ReLU(inplace=True), - nn.Linear(512, num_heads, bias=False)) + self.cpb_mlp = nn.Sequential( + nn.Linear(2, 512, bias=True), + nn.ReLU(inplace=True), + nn.Linear(512, num_heads, bias=False), + ) # get relative_coords_table - relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32) - relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32) - relative_coords_table = torch.stack( - torch.meshgrid([relative_coords_h, - relative_coords_w])).permute(1, 2, 0).contiguous().unsqueeze(0) # 1, 2*Wh-1, 2*Ww-1, 2 + relative_coords_h = torch.arange( + -(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32 + ) + relative_coords_w = torch.arange( + -(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32 + ) + relative_coords_table = ( + torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w])) + .permute(1, 2, 0) + .contiguous() + .unsqueeze(0) + ) # 1, 2*Wh-1, 2*Ww-1, 2 if pretrained_window_size[0] > 0: - relative_coords_table[:, :, :, 0] /= (pretrained_window_size[0] - 1) - relative_coords_table[:, :, :, 1] /= (pretrained_window_size[1] - 1) + relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1 + relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1 else: - relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1) - relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1) + relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1 + relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1 relative_coords_table *= 8 # normalize to -8, 8 - relative_coords_table = torch.sign(relative_coords_table) * torch.log2( - torch.abs(relative_coords_table) + 1.0) / np.log2(8) + relative_coords_table = ( + torch.sign(relative_coords_table) + * torch.log2(torch.abs(relative_coords_table) + 1.0) + / np.log2(8) + ) self.register_buffer("relative_coords_table", relative_coords_table) @@ -188,8 +245,12 @@ def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, at coords_w = torch.arange(self.window_size[1]) coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww - relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww - relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords = ( + coords_flatten[:, :, None] - coords_flatten[:, None, :] + ) # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute( + 1, 2, 0 + ).contiguous() # Wh*Ww, Wh*Ww, 2 relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 relative_coords[:, :, 1] += self.window_size[1] - 1 relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 @@ -209,21 +270,40 @@ def forward(self, x, mask=None): mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None """ B_, N, C = x.shape - qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) - q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + qkv = ( + self.qkv(x) + .reshape(B_, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + q, k, v = ( + qkv[0], + qkv[1], + qkv[2], + ) # make torchscript happy (cannot use tensor as tuple) q = q * self.scale - attn = (q @ k.transpose(-2, -1)) + attn = q @ k.transpose(-2, -1) - relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads) - relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view( - self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH - relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view( + -1, self.num_heads + ) + relative_position_bias = relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1], + self.window_size[0] * self.window_size[1], + -1, + ) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() # nH, Wh*Ww, Wh*Ww attn = attn + relative_position_bias.unsqueeze(0) if mask is not None: nW = mask.shape[0] - attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze( + 1 + ).unsqueeze(0) attn = attn.view(-1, self.num_heads, N, N) attn = self.softmax(attn) else: @@ -237,8 +317,10 @@ def forward(self, x, mask=None): return x def extra_repr(self) -> str: - return f'dim={self.dim}, window_size={self.window_size}, ' \ - f'pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}' + return ( + f"dim={self.dim}, window_size={self.window_size}, " + f"pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}" + ) def flops(self, N): # calculate flops for 1 window with token length of N @@ -255,7 +337,7 @@ def flops(self, N): class SwinTransformerBlock(nn.Module): - r""" Swin Transformer Block. + r"""Swin Transformer Block. Args: dim (int): Number of input channels. @@ -289,12 +371,37 @@ class SwinTransformerBlock(nn.Module): moe_drop (float): Dropout rate in MoE. Default: 0.0 """ - def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., - act_layer=nn.GELU, norm_layer=nn.LayerNorm, mlp_fc2_bias=True, init_std=0.02, pretrained_window_size=0, - is_moe=False, num_local_experts=1, top_value=1, capacity_factor=1.25, cosine_router=False, - normalize_gate=False, use_bpr=True, is_gshard_loss=True, gate_noise=1.0, - cosine_router_dim=256, cosine_router_init_t=0.5, moe_drop=0.0): + def __init__( + self, + dim, + input_resolution, + num_heads, + window_size=7, + shift_size=0, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + mlp_fc2_bias=True, + init_std=0.02, + pretrained_window_size=0, + is_moe=False, + num_local_experts=1, + top_value=1, + capacity_factor=1.25, + cosine_router=False, + normalize_gate=False, + use_bpr=True, + is_gshard_loss=True, + gate_noise=1.0, + cosine_router_dim=256, + cosine_router_init_t=0.5, + moe_drop=0.0, + ): super().__init__() self.dim = dim self.input_resolution = input_resolution @@ -310,57 +417,80 @@ def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0 # if window size is larger than input resolution, we don't partition windows self.shift_size = 0 self.window_size = min(self.input_resolution) - assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + assert ( + 0 <= self.shift_size < self.window_size + ), "shift_size must in 0-window_size" self.norm1 = norm_layer(dim) self.attn = WindowAttention( - dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, - qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, - pretrained_window_size=to_2tuple(pretrained_window_size)) + dim, + window_size=to_2tuple(self.window_size), + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + pretrained_window_size=to_2tuple(pretrained_window_size), + ) - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) if self.is_moe: - self.mlp = MoEMlp(in_features=dim, - hidden_features=mlp_hidden_dim, - num_local_experts=num_local_experts, - top_value=top_value, - capacity_factor=capacity_factor, - cosine_router=cosine_router, - normalize_gate=normalize_gate, - use_bpr=use_bpr, - is_gshard_loss=is_gshard_loss, - gate_noise=gate_noise, - cosine_router_dim=cosine_router_dim, - cosine_router_init_t=cosine_router_init_t, - moe_drop=moe_drop, - mlp_fc2_bias=mlp_fc2_bias, - init_std=init_std) + self.mlp = MoEMlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + num_local_experts=num_local_experts, + top_value=top_value, + capacity_factor=capacity_factor, + cosine_router=cosine_router, + normalize_gate=normalize_gate, + use_bpr=use_bpr, + is_gshard_loss=is_gshard_loss, + gate_noise=gate_noise, + cosine_router_dim=cosine_router_dim, + cosine_router_init_t=cosine_router_init_t, + moe_drop=moe_drop, + mlp_fc2_bias=mlp_fc2_bias, + init_std=init_std, + ) else: - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop, - mlp_fc2_bias=mlp_fc2_bias) + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + mlp_fc2_bias=mlp_fc2_bias, + ) if self.shift_size > 0: # calculate attention mask for SW-MSA H, W = self.input_resolution img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 - h_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) - w_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) + h_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + w_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 - mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = window_partition( + img_mask, self.window_size + ) # nW, window_size, window_size, 1 mask_windows = mask_windows.view(-1, self.window_size * self.window_size) attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + attn_mask = attn_mask.masked_fill( + attn_mask != 0, float(-100.0) + ).masked_fill(attn_mask == 0, float(0.0)) else: attn_mask = None @@ -377,16 +507,24 @@ def forward(self, x): # cyclic shift if self.shift_size > 0: - shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + shifted_x = torch.roll( + x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2) + ) else: shifted_x = x # partition windows - x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C - x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + x_windows = window_partition( + shifted_x, self.window_size + ) # nW*B, window_size, window_size, C + x_windows = x_windows.view( + -1, self.window_size * self.window_size, C + ) # nW*B, window_size*window_size, C # W-MSA/SW-MSA - attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C + attn_windows = self.attn( + x_windows, mask=self.attn_mask + ) # nW*B, window_size*window_size, C # merge windows attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) @@ -394,7 +532,9 @@ def forward(self, x): # reverse cyclic shift if self.shift_size > 0: - x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + x = torch.roll( + shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2) + ) else: x = shifted_x x = x.view(B, H * W, C) @@ -412,8 +552,10 @@ def forward(self, x): return x def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ - f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + return ( + f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " + f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + ) def flops(self): flops = 0 @@ -425,7 +567,16 @@ def flops(self): flops += nW * self.attn.flops(self.window_size * self.window_size) # mlp if self.is_moe: - flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio * self.capacity_factor * self.top_value + flops += ( + 2 + * H + * W + * self.dim + * self.dim + * self.mlp_ratio + * self.capacity_factor + * self.top_value + ) else: flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio # norm2 @@ -434,7 +585,7 @@ def flops(self): class PatchMerging(nn.Module): - r""" Patch Merging Layer. + r"""Patch Merging Layer. Args: input_resolution (tuple[int]): Resolution of input feature. @@ -483,7 +634,7 @@ def flops(self): class BasicLayer(nn.Module): - """ A basic Swin Transformer layer for one stage. + """A basic Swin Transformer layer for one stage. Args: dim (int): Number of input channels. @@ -518,13 +669,38 @@ class BasicLayer(nn.Module): moe_drop (float): Dropout rate in MoE. Default: 0.0 """ - def __init__(self, dim, input_resolution, depth, num_heads, window_size, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., norm_layer=nn.LayerNorm, downsample=None, - mlp_fc2_bias=True, init_std=0.02, use_checkpoint=False, pretrained_window_size=0, - moe_block=[-1], num_local_experts=1, top_value=1, capacity_factor=1.25, cosine_router=False, - normalize_gate=False, use_bpr=True, is_gshard_loss=True, - cosine_router_dim=256, cosine_router_init_t=0.5, gate_noise=1.0, moe_drop=0.0): + def __init__( + self, + dim, + input_resolution, + depth, + num_heads, + window_size, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + norm_layer=nn.LayerNorm, + downsample=None, + mlp_fc2_bias=True, + init_std=0.02, + use_checkpoint=False, + pretrained_window_size=0, + moe_block=[-1], + num_local_experts=1, + top_value=1, + capacity_factor=1.25, + cosine_router=False, + normalize_gate=False, + use_bpr=True, + is_gshard_loss=True, + cosine_router_dim=256, + cosine_router_init_t=0.5, + gate_noise=1.0, + moe_drop=0.0, + ): super().__init__() self.dim = dim @@ -533,36 +709,48 @@ def __init__(self, dim, input_resolution, depth, num_heads, window_size, self.use_checkpoint = use_checkpoint # build blocks - self.blocks = nn.ModuleList([ - SwinTransformerBlock(dim=dim, input_resolution=input_resolution, - num_heads=num_heads, window_size=window_size, - shift_size=0 if (i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop, attn_drop=attn_drop, - drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, - norm_layer=norm_layer, - mlp_fc2_bias=mlp_fc2_bias, - init_std=init_std, - pretrained_window_size=pretrained_window_size, - - is_moe=True if i in moe_block else False, - num_local_experts=num_local_experts, - top_value=top_value, - capacity_factor=capacity_factor, - cosine_router=cosine_router, - normalize_gate=normalize_gate, - use_bpr=use_bpr, - is_gshard_loss=is_gshard_loss, - gate_noise=gate_noise, - cosine_router_dim=cosine_router_dim, - cosine_router_init_t=cosine_router_init_t, - moe_drop=moe_drop) - for i in range(depth)]) + self.blocks = nn.ModuleList( + [ + SwinTransformerBlock( + dim=dim, + input_resolution=input_resolution, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] + if isinstance(drop_path, list) + else drop_path, + norm_layer=norm_layer, + mlp_fc2_bias=mlp_fc2_bias, + init_std=init_std, + pretrained_window_size=pretrained_window_size, + is_moe=True if i in moe_block else False, + num_local_experts=num_local_experts, + top_value=top_value, + capacity_factor=capacity_factor, + cosine_router=cosine_router, + normalize_gate=normalize_gate, + use_bpr=use_bpr, + is_gshard_loss=is_gshard_loss, + gate_noise=gate_noise, + cosine_router_dim=cosine_router_dim, + cosine_router_init_t=cosine_router_init_t, + moe_drop=moe_drop, + ) + for i in range(depth) + ] + ) # patch merging layer if downsample is not None: - self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer) + self.downsample = downsample( + input_resolution, dim=dim, norm_layer=norm_layer + ) else: self.downsample = None @@ -597,7 +785,7 @@ def flops(self): class PatchEmbed(nn.Module): - r""" Image to Patch Embedding + r"""Image to Patch Embedding Args: img_size (int): Image size. Default: 224. @@ -607,11 +795,16 @@ class PatchEmbed(nn.Module): norm_layer (nn.Module, optional): Normalization layer. Default: None """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + def __init__( + self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None + ): super().__init__() img_size = to_2tuple(img_size) patch_size = to_2tuple(patch_size) - patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] + patches_resolution = [ + img_size[0] // patch_size[0], + img_size[1] // patch_size[1], + ] self.img_size = img_size self.patch_size = patch_size self.patches_resolution = patches_resolution @@ -620,7 +813,9 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_la self.in_chans = in_chans self.embed_dim = embed_dim - self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=patch_size, stride=patch_size + ) if norm_layer is not None: self.norm = norm_layer(embed_dim) else: @@ -629,8 +824,9 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_la def forward(self, x): B, C, H, W = x.shape # FIXME look at relaxing size constraints - assert H == self.img_size[0] and W == self.img_size[1], \ - f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." + assert ( + H == self.img_size[0] and W == self.img_size[1] + ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C if self.norm is not None: x = self.norm(x) @@ -638,14 +834,20 @@ def forward(self, x): def flops(self): Ho, Wo = self.patches_resolution - flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + flops = ( + Ho + * Wo + * self.embed_dim + * self.in_chans + * (self.patch_size[0] * self.patch_size[1]) + ) if self.norm is not None: flops += Ho * Wo * self.embed_dim return flops class SwinTransformerMoE(nn.Module): - r""" Swin Transformer + r"""Swin Transformer A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - https://arxiv.org/pdf/2103.14030 @@ -687,15 +889,44 @@ class SwinTransformerMoE(nn.Module): aux_loss_weight (float): auxiliary loss weight. Default: 0.1 """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, - embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], - window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None, - drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1, - norm_layer=nn.LayerNorm, ape=False, patch_norm=True, - mlp_fc2_bias=True, init_std=0.02, use_checkpoint=False, pretrained_window_sizes=[0, 0, 0, 0], - moe_blocks=[[-1], [-1], [-1], [-1]], num_local_experts=1, top_value=1, capacity_factor=1.25, - cosine_router=False, normalize_gate=False, use_bpr=True, is_gshard_loss=True, gate_noise=1.0, - cosine_router_dim=256, cosine_router_init_t=0.5, moe_drop=0.0, aux_loss_weight=0.01, **kwargs): + def __init__( + self, + img_size=224, + patch_size=4, + in_chans=3, + num_classes=1000, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.1, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + mlp_fc2_bias=True, + init_std=0.02, + use_checkpoint=False, + pretrained_window_sizes=[0, 0, 0, 0], + moe_blocks=[[-1], [-1], [-1], [-1]], + num_local_experts=1, + top_value=1, + capacity_factor=1.25, + cosine_router=False, + normalize_gate=False, + use_bpr=True, + is_gshard_loss=True, + gate_noise=1.0, + cosine_router_dim=256, + cosine_router_init_t=0.5, + moe_drop=0.0, + aux_loss_weight=0.01, + **kwargs, + ): super().__init__() self._ddp_params_and_buffers_to_ignore = list() @@ -709,65 +940,87 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, self.init_std = init_std self.aux_loss_weight = aux_loss_weight self.num_local_experts = num_local_experts - self.global_experts = num_local_experts * dist.get_world_size() if num_local_experts > 0 \ + self.global_experts = ( + num_local_experts * dist.get_world_size() + if num_local_experts > 0 else dist.get_world_size() // (-num_local_experts) - self.sharded_count = (1.0 / num_local_experts) if num_local_experts > 0 else (-num_local_experts) + ) + self.sharded_count = ( + (1.0 / num_local_experts) if num_local_experts > 0 else (-num_local_experts) + ) # split image into non-overlapping patches self.patch_embed = PatchEmbed( - img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None) + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, + ) num_patches = self.patch_embed.num_patches patches_resolution = self.patch_embed.patches_resolution self.patches_resolution = patches_resolution # absolute position embedding if self.ape: - self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) + self.absolute_pos_embed = nn.Parameter( + torch.zeros(1, num_patches, embed_dim) + ) trunc_normal_(self.absolute_pos_embed, std=self.init_std) self.pos_drop = nn.Dropout(p=drop_rate) # stochastic depth - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule # build layers self.layers = nn.ModuleList() for i_layer in range(self.num_layers): - layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer), - input_resolution=(patches_resolution[0] // (2 ** i_layer), - patches_resolution[1] // (2 ** i_layer)), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, - drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], - norm_layer=norm_layer, - downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, - mlp_fc2_bias=mlp_fc2_bias, - init_std=init_std, - use_checkpoint=use_checkpoint, - pretrained_window_size=pretrained_window_sizes[i_layer], - - moe_block=moe_blocks[i_layer], - num_local_experts=num_local_experts, - top_value=top_value, - capacity_factor=capacity_factor, - cosine_router=cosine_router, - normalize_gate=normalize_gate, - use_bpr=use_bpr, - is_gshard_loss=is_gshard_loss, - gate_noise=gate_noise, - cosine_router_dim=cosine_router_dim, - cosine_router_init_t=cosine_router_init_t, - moe_drop=moe_drop) + layer = BasicLayer( + dim=int(embed_dim * 2**i_layer), + input_resolution=( + patches_resolution[0] // (2**i_layer), + patches_resolution[1] // (2**i_layer), + ), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=self.mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + mlp_fc2_bias=mlp_fc2_bias, + init_std=init_std, + use_checkpoint=use_checkpoint, + pretrained_window_size=pretrained_window_sizes[i_layer], + moe_block=moe_blocks[i_layer], + num_local_experts=num_local_experts, + top_value=top_value, + capacity_factor=capacity_factor, + cosine_router=cosine_router, + normalize_gate=normalize_gate, + use_bpr=use_bpr, + is_gshard_loss=is_gshard_loss, + gate_noise=gate_noise, + cosine_router_dim=cosine_router_dim, + cosine_router_init_t=cosine_router_init_t, + moe_drop=moe_drop, + ) self.layers.append(layer) self.norm = norm_layer(self.num_features) self.avgpool = nn.AdaptiveAvgPool1d(1) - self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity() + self.head = ( + nn.Linear(self.num_features, num_classes) + if num_classes > 0 + else nn.Identity() + ) self.apply(self._init_weights) @@ -784,12 +1037,19 @@ def _init_weights(self, m): @torch.jit.ignore def no_weight_decay(self): - return {'absolute_pos_embed'} + return {"absolute_pos_embed"} @torch.jit.ignore def no_weight_decay_keywords(self): - return {"cpb_mlp", 'relative_position_bias_table', 'fc1_bias', 'fc2_bias', - 'temperature', 'cosine_projector', 'sim_matrix'} + return { + "cpb_mlp", + "relative_position_bias_table", + "fc1_bias", + "fc2_bias", + "temperature", + "cosine_projector", + "sim_matrix", + } def forward_features(self, x): x = self.patch_embed(x) @@ -819,6 +1079,11 @@ def flops(self): flops += self.patch_embed.flops() for i, layer in enumerate(self.layers): flops += layer.flops() - flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers) + flops += ( + self.num_features + * self.patches_resolution[0] + * self.patches_resolution[1] + // (2**self.num_layers) + ) flops += self.num_features * self.num_classes return flops diff --git a/models/swin_transformer_v2.py b/models/swin_transformer_v2.py index a429d0a2..9d6d2f86 100644 --- a/models/swin_transformer_v2.py +++ b/models/swin_transformer_v2.py @@ -5,16 +5,23 @@ # Written by Ze Liu # -------------------------------------------------------- +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint as checkpoint from timm.models.layers import DropPath, to_2tuple, trunc_normal_ -import numpy as np class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + drop=0.0, + ): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -43,7 +50,9 @@ def window_partition(x, window_size): """ B, H, W, C = x.shape x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) - windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + windows = ( + x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + ) return windows @@ -59,13 +68,15 @@ def window_reverse(windows, window_size, H, W): x: (B, H, W, C) """ B = int(windows.shape[0] / (H * W / window_size / window_size)) - x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = windows.view( + B, H // window_size, W // window_size, window_size, window_size, -1 + ) x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) return x class WindowAttention(nn.Module): - r""" Window based multi-head self attention (W-MSA) module with relative position bias. + r"""Window based multi-head self attention (W-MSA) module with relative position bias. It supports both of shifted and non-shifted window. Args: @@ -78,8 +89,16 @@ class WindowAttention(nn.Module): pretrained_window_size (tuple[int]): The height and width of the window in pre-training. """ - def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0., - pretrained_window_size=[0, 0]): + def __init__( + self, + dim, + window_size, + num_heads, + qkv_bias=True, + attn_drop=0.0, + proj_drop=0.0, + pretrained_window_size=[0, 0], + ): super().__init__() self.dim = dim @@ -87,28 +106,42 @@ def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., pro self.pretrained_window_size = pretrained_window_size self.num_heads = num_heads - self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True) + self.logit_scale = nn.Parameter( + torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True + ) # mlp to generate continuous relative position bias - self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True), - nn.ReLU(inplace=True), - nn.Linear(512, num_heads, bias=False)) + self.cpb_mlp = nn.Sequential( + nn.Linear(2, 512, bias=True), + nn.ReLU(inplace=True), + nn.Linear(512, num_heads, bias=False), + ) # get relative_coords_table - relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32) - relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32) - relative_coords_table = torch.stack( - torch.meshgrid([relative_coords_h, - relative_coords_w])).permute(1, 2, 0).contiguous().unsqueeze(0) # 1, 2*Wh-1, 2*Ww-1, 2 + relative_coords_h = torch.arange( + -(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32 + ) + relative_coords_w = torch.arange( + -(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32 + ) + relative_coords_table = ( + torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w])) + .permute(1, 2, 0) + .contiguous() + .unsqueeze(0) + ) # 1, 2*Wh-1, 2*Ww-1, 2 if pretrained_window_size[0] > 0: - relative_coords_table[:, :, :, 0] /= (pretrained_window_size[0] - 1) - relative_coords_table[:, :, :, 1] /= (pretrained_window_size[1] - 1) + relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1 + relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1 else: - relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1) - relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1) + relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1 + relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1 relative_coords_table *= 8 # normalize to -8, 8 - relative_coords_table = torch.sign(relative_coords_table) * torch.log2( - torch.abs(relative_coords_table) + 1.0) / np.log2(8) + relative_coords_table = ( + torch.sign(relative_coords_table) + * torch.log2(torch.abs(relative_coords_table) + 1.0) + / np.log2(8) + ) self.register_buffer("relative_coords_table", relative_coords_table) @@ -117,8 +150,12 @@ def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., pro coords_w = torch.arange(self.window_size[1]) coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww - relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww - relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords = ( + coords_flatten[:, :, None] - coords_flatten[:, None, :] + ) # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute( + 1, 2, 0 + ).contiguous() # Wh*Ww, Wh*Ww, 2 relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 relative_coords[:, :, 1] += self.window_size[1] - 1 relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 @@ -146,26 +183,49 @@ def forward(self, x, mask=None): B_, N, C = x.shape qkv_bias = None if self.q_bias is not None: - qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias)) + qkv_bias = torch.cat( + ( + self.q_bias, + torch.zeros_like(self.v_bias, requires_grad=False), + self.v_bias, + ) + ) qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias) qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) - q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + q, k, v = ( + qkv[0], + qkv[1], + qkv[2], + ) # make torchscript happy (cannot use tensor as tuple) # cosine attention - attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)) - logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / 0.01))).exp() + attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) + logit_scale = torch.clamp( + self.logit_scale, max=torch.log(torch.tensor(1.0 / 0.01)) + ).exp() attn = attn * logit_scale - relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads) - relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view( - self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH - relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view( + -1, self.num_heads + ) + relative_position_bias = relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1], + self.window_size[0] * self.window_size[1], + -1, + ) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() # nH, Wh*Ww, Wh*Ww relative_position_bias = 16 * torch.sigmoid(relative_position_bias) attn = attn + relative_position_bias.unsqueeze(0) if mask is not None: nW = mask.shape[0] - attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze( + 1 + ).unsqueeze(0) attn = attn.view(-1, self.num_heads, N, N) attn = self.softmax(attn) else: @@ -179,8 +239,10 @@ def forward(self, x, mask=None): return x def extra_repr(self) -> str: - return f'dim={self.dim}, window_size={self.window_size}, ' \ - f'pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}' + return ( + f"dim={self.dim}, window_size={self.window_size}, " + f"pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}" + ) def flops(self, N): # calculate flops for 1 window with token length of N @@ -197,7 +259,7 @@ def flops(self, N): class SwinTransformerBlock(nn.Module): - r""" Swin Transformer Block. + r"""Swin Transformer Block. Args: dim (int): Number of input channels. @@ -215,9 +277,22 @@ class SwinTransformerBlock(nn.Module): pretrained_window_size (int): Window size in pre-training. """ - def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, - mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0., - act_layer=nn.GELU, norm_layer=nn.LayerNorm, pretrained_window_size=0): + def __init__( + self, + dim, + input_resolution, + num_heads, + window_size=7, + shift_size=0, + mlp_ratio=4.0, + qkv_bias=True, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + pretrained_window_size=0, + ): super().__init__() self.dim = dim self.input_resolution = input_resolution @@ -229,39 +304,59 @@ def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0 # if window size is larger than input resolution, we don't partition windows self.shift_size = 0 self.window_size = min(self.input_resolution) - assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + assert ( + 0 <= self.shift_size < self.window_size + ), "shift_size must in 0-window_size" self.norm1 = norm_layer(dim) self.attn = WindowAttention( - dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, - qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, - pretrained_window_size=to_2tuple(pretrained_window_size)) - - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + dim, + window_size=to_2tuple(self.window_size), + num_heads=num_heads, + qkv_bias=qkv_bias, + attn_drop=attn_drop, + proj_drop=drop, + pretrained_window_size=to_2tuple(pretrained_window_size), + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + ) if self.shift_size > 0: # calculate attention mask for SW-MSA H, W = self.input_resolution img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 - h_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) - w_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) + h_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + w_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 - mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = window_partition( + img_mask, self.window_size + ) # nW, window_size, window_size, 1 mask_windows = mask_windows.view(-1, self.window_size * self.window_size) attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + attn_mask = attn_mask.masked_fill( + attn_mask != 0, float(-100.0) + ).masked_fill(attn_mask == 0, float(0.0)) else: attn_mask = None @@ -277,16 +372,24 @@ def forward(self, x): # cyclic shift if self.shift_size > 0: - shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + shifted_x = torch.roll( + x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2) + ) else: shifted_x = x # partition windows - x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C - x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + x_windows = window_partition( + shifted_x, self.window_size + ) # nW*B, window_size, window_size, C + x_windows = x_windows.view( + -1, self.window_size * self.window_size, C + ) # nW*B, window_size*window_size, C # W-MSA/SW-MSA - attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C + attn_windows = self.attn( + x_windows, mask=self.attn_mask + ) # nW*B, window_size*window_size, C # merge windows attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) @@ -294,7 +397,9 @@ def forward(self, x): # reverse cyclic shift if self.shift_size > 0: - x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + x = torch.roll( + shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2) + ) else: x = shifted_x x = x.view(B, H * W, C) @@ -306,8 +411,10 @@ def forward(self, x): return x def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ - f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + return ( + f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " + f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + ) def flops(self): flops = 0 @@ -325,7 +432,7 @@ def flops(self): class PatchMerging(nn.Module): - r""" Patch Merging Layer. + r"""Patch Merging Layer. Args: input_resolution (tuple[int]): Resolution of input feature. @@ -374,7 +481,7 @@ def flops(self): class BasicLayer(nn.Module): - """ A basic Swin Transformer layer for one stage. + """A basic Swin Transformer layer for one stage. Args: dim (int): Number of input channels. @@ -393,10 +500,23 @@ class BasicLayer(nn.Module): pretrained_window_size (int): Local window size in pre-training. """ - def __init__(self, dim, input_resolution, depth, num_heads, window_size, - mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., - drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False, - pretrained_window_size=0): + def __init__( + self, + dim, + input_resolution, + depth, + num_heads, + window_size, + mlp_ratio=4.0, + qkv_bias=True, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False, + pretrained_window_size=0, + ): super().__init__() self.dim = dim @@ -405,21 +525,33 @@ def __init__(self, dim, input_resolution, depth, num_heads, window_size, self.use_checkpoint = use_checkpoint # build blocks - self.blocks = nn.ModuleList([ - SwinTransformerBlock(dim=dim, input_resolution=input_resolution, - num_heads=num_heads, window_size=window_size, - shift_size=0 if (i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - drop=drop, attn_drop=attn_drop, - drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, - norm_layer=norm_layer, - pretrained_window_size=pretrained_window_size) - for i in range(depth)]) + self.blocks = nn.ModuleList( + [ + SwinTransformerBlock( + dim=dim, + input_resolution=input_resolution, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] + if isinstance(drop_path, list) + else drop_path, + norm_layer=norm_layer, + pretrained_window_size=pretrained_window_size, + ) + for i in range(depth) + ] + ) # patch merging layer if downsample is not None: - self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer) + self.downsample = downsample( + input_resolution, dim=dim, norm_layer=norm_layer + ) else: self.downsample = None @@ -453,7 +585,7 @@ def _init_respostnorm(self): class PatchEmbed(nn.Module): - r""" Image to Patch Embedding + r"""Image to Patch Embedding Args: img_size (int): Image size. Default: 224. @@ -463,11 +595,16 @@ class PatchEmbed(nn.Module): norm_layer (nn.Module, optional): Normalization layer. Default: None """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + def __init__( + self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None + ): super().__init__() img_size = to_2tuple(img_size) patch_size = to_2tuple(patch_size) - patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] + patches_resolution = [ + img_size[0] // patch_size[0], + img_size[1] // patch_size[1], + ] self.img_size = img_size self.patch_size = patch_size self.patches_resolution = patches_resolution @@ -476,7 +613,9 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_la self.in_chans = in_chans self.embed_dim = embed_dim - self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=patch_size, stride=patch_size + ) if norm_layer is not None: self.norm = norm_layer(embed_dim) else: @@ -485,8 +624,9 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_la def forward(self, x): B, C, H, W = x.shape # FIXME look at relaxing size constraints - assert H == self.img_size[0] and W == self.img_size[1], \ - f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." + assert ( + H == self.img_size[0] and W == self.img_size[1] + ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C if self.norm is not None: x = self.norm(x) @@ -494,14 +634,20 @@ def forward(self, x): def flops(self): Ho, Wo = self.patches_resolution - flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + flops = ( + Ho + * Wo + * self.embed_dim + * self.in_chans + * (self.patch_size[0] * self.patch_size[1]) + ) if self.norm is not None: flops += Ho * Wo * self.embed_dim return flops class SwinTransformerV2(nn.Module): - r""" Swin Transformer + r"""Swin Transformer A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - https://arxiv.org/pdf/2103.14030 @@ -526,12 +672,28 @@ class SwinTransformerV2(nn.Module): pretrained_window_sizes (tuple(int)): Pretrained window sizes of each layer. """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, - embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], - window_size=7, mlp_ratio=4., qkv_bias=True, - drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1, - norm_layer=nn.LayerNorm, ape=False, patch_norm=True, - use_checkpoint=False, pretrained_window_sizes=[0, 0, 0, 0], **kwargs): + def __init__( + self, + img_size=224, + patch_size=4, + in_chans=3, + num_classes=1000, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4.0, + qkv_bias=True, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.1, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + use_checkpoint=False, + pretrained_window_sizes=[0, 0, 0, 0], + **kwargs, + ): super().__init__() self.num_classes = num_classes @@ -544,44 +706,61 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, # split image into non-overlapping patches self.patch_embed = PatchEmbed( - img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None) + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, + ) num_patches = self.patch_embed.num_patches patches_resolution = self.patch_embed.patches_resolution self.patches_resolution = patches_resolution # absolute position embedding if self.ape: - self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) - trunc_normal_(self.absolute_pos_embed, std=.02) + self.absolute_pos_embed = nn.Parameter( + torch.zeros(1, num_patches, embed_dim) + ) + trunc_normal_(self.absolute_pos_embed, std=0.02) self.pos_drop = nn.Dropout(p=drop_rate) # stochastic depth - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule # build layers self.layers = nn.ModuleList() for i_layer in range(self.num_layers): - layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer), - input_resolution=(patches_resolution[0] // (2 ** i_layer), - patches_resolution[1] // (2 ** i_layer)), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, - drop=drop_rate, attn_drop=attn_drop_rate, - drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], - norm_layer=norm_layer, - downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, - use_checkpoint=use_checkpoint, - pretrained_window_size=pretrained_window_sizes[i_layer]) + layer = BasicLayer( + dim=int(embed_dim * 2**i_layer), + input_resolution=( + patches_resolution[0] // (2**i_layer), + patches_resolution[1] // (2**i_layer), + ), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=self.mlp_ratio, + qkv_bias=qkv_bias, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint, + pretrained_window_size=pretrained_window_sizes[i_layer], + ) self.layers.append(layer) self.norm = norm_layer(self.num_features) self.avgpool = nn.AdaptiveAvgPool1d(1) - self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity() + self.head = ( + nn.Linear(self.num_features, num_classes) + if num_classes > 0 + else nn.Identity() + ) self.apply(self._init_weights) for bly in self.layers: @@ -589,7 +768,7 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -598,11 +777,11 @@ def _init_weights(self, m): @torch.jit.ignore def no_weight_decay(self): - return {'absolute_pos_embed'} + return {"absolute_pos_embed"} @torch.jit.ignore def no_weight_decay_keywords(self): - return {"cpb_mlp", "logit_scale", 'relative_position_bias_table'} + return {"cpb_mlp", "logit_scale", "relative_position_bias_table"} def forward_features(self, x): x = self.patch_embed(x) @@ -628,6 +807,11 @@ def flops(self): flops += self.patch_embed.flops() for i, layer in enumerate(self.layers): flops += layer.flops() - flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers) + flops += ( + self.num_features + * self.patches_resolution[0] + * self.patches_resolution[1] + // (2**self.num_layers) + ) flops += self.num_features * self.num_classes return flops diff --git a/optimizer.py b/optimizer.py index 44317019..fc1f3167 100644 --- a/optimizer.py +++ b/optimizer.py @@ -6,6 +6,7 @@ # -------------------------------------------------------- from functools import partial + from torch import optim as optim try: @@ -22,36 +23,72 @@ def build_optimizer(config, model, simmim=False, is_pretrain=False): """ skip = {} skip_keywords = {} - if hasattr(model, 'no_weight_decay'): + if hasattr(model, "no_weight_decay"): skip = model.no_weight_decay() - if hasattr(model, 'no_weight_decay_keywords'): + if hasattr(model, "no_weight_decay_keywords"): skip_keywords = model.no_weight_decay_keywords() if simmim: if is_pretrain: parameters = get_pretrain_param_groups(model, skip, skip_keywords) else: - depths = config.MODEL.SWIN.DEPTHS if config.MODEL.TYPE == 'swin' else config.MODEL.SWINV2.DEPTHS + depths = ( + config.MODEL.SWIN.DEPTHS + if config.MODEL.TYPE == "swin" + else config.MODEL.SWINV2.DEPTHS + ) num_layers = sum(depths) - get_layer_func = partial(get_swin_layer, num_layers=num_layers + 2, depths=depths) - scales = list(config.TRAIN.LAYER_DECAY ** i for i in reversed(range(num_layers + 2))) - parameters = get_finetune_param_groups(model, config.TRAIN.BASE_LR, config.TRAIN.WEIGHT_DECAY, get_layer_func, scales, skip, skip_keywords) + get_layer_func = partial( + get_swin_layer, num_layers=num_layers + 2, depths=depths + ) + scales = list( + config.TRAIN.LAYER_DECAY**i for i in reversed(range(num_layers + 2)) + ) + parameters = get_finetune_param_groups( + model, + config.TRAIN.BASE_LR, + config.TRAIN.WEIGHT_DECAY, + get_layer_func, + scales, + skip, + skip_keywords, + ) else: parameters = set_weight_decay(model, skip, skip_keywords) opt_lower = config.TRAIN.OPTIMIZER.NAME.lower() optimizer = None - if opt_lower == 'sgd': - optimizer = optim.SGD(parameters, momentum=config.TRAIN.OPTIMIZER.MOMENTUM, nesterov=True, - lr=config.TRAIN.BASE_LR, weight_decay=config.TRAIN.WEIGHT_DECAY) - elif opt_lower == 'adamw': - optimizer = optim.AdamW(parameters, eps=config.TRAIN.OPTIMIZER.EPS, betas=config.TRAIN.OPTIMIZER.BETAS, - lr=config.TRAIN.BASE_LR, weight_decay=config.TRAIN.WEIGHT_DECAY) - elif opt_lower == 'fused_adam': - optimizer = FusedAdam(parameters, eps=config.TRAIN.OPTIMIZER.EPS, betas=config.TRAIN.OPTIMIZER.BETAS, - lr=config.TRAIN.BASE_LR, weight_decay=config.TRAIN.WEIGHT_DECAY) - elif opt_lower == 'fused_lamb': - optimizer = FusedLAMB(parameters, eps=config.TRAIN.OPTIMIZER.EPS, betas=config.TRAIN.OPTIMIZER.BETAS, - lr=config.TRAIN.BASE_LR, weight_decay=config.TRAIN.WEIGHT_DECAY) + if opt_lower == "sgd": + optimizer = optim.SGD( + parameters, + momentum=config.TRAIN.OPTIMIZER.MOMENTUM, + nesterov=True, + lr=config.TRAIN.BASE_LR, + weight_decay=config.TRAIN.WEIGHT_DECAY, + ) + elif opt_lower == "adamw": + optimizer = optim.AdamW( + parameters, + eps=config.TRAIN.OPTIMIZER.EPS, + betas=config.TRAIN.OPTIMIZER.BETAS, + lr=config.TRAIN.BASE_LR, + weight_decay=config.TRAIN.WEIGHT_DECAY, + ) + elif opt_lower == "fused_adam": + optimizer = FusedAdam( + parameters, + eps=config.TRAIN.OPTIMIZER.EPS, + betas=config.TRAIN.OPTIMIZER.BETAS, + lr=config.TRAIN.BASE_LR, + weight_decay=config.TRAIN.WEIGHT_DECAY, + ) + elif opt_lower == "fused_lamb": + optimizer = FusedLAMB( + parameters, + eps=config.TRAIN.OPTIMIZER.EPS, + betas=config.TRAIN.OPTIMIZER.BETAS, + lr=config.TRAIN.BASE_LR, + weight_decay=config.TRAIN.WEIGHT_DECAY, + ) return optimizer @@ -63,14 +100,17 @@ def set_weight_decay(model, skip_list=(), skip_keywords=()): for name, param in model.named_parameters(): if not param.requires_grad: continue # frozen weights - if len(param.shape) == 1 or name.endswith(".bias") or (name in skip_list) or \ - check_keywords_in_name(name, skip_keywords): + if ( + len(param.shape) == 1 + or name.endswith(".bias") + or (name in skip_list) + or check_keywords_in_name(name, skip_keywords) + ): no_decay.append(param) # print(f"{name} has no weight decay") else: has_decay.append(param) - return [{'params': has_decay}, - {'params': no_decay, 'weight_decay': 0.}] + return [{"params": has_decay}, {"params": no_decay, "weight_decay": 0.0}] def check_keywords_in_name(name, keywords=()): @@ -86,19 +126,22 @@ def get_pretrain_param_groups(model, skip_list=(), skip_keywords=()): no_decay = [] has_decay_name = [] no_decay_name = [] - + for name, param in model.named_parameters(): if not param.requires_grad: continue - if len(param.shape) == 1 or name.endswith(".bias") or (name in skip_list) or \ - check_keywords_in_name(name, skip_keywords): + if ( + len(param.shape) == 1 + or name.endswith(".bias") + or (name in skip_list) + or check_keywords_in_name(name, skip_keywords) + ): no_decay.append(param) no_decay_name.append(name) else: has_decay.append(param) has_decay_name.append(name) - return [{'params': has_decay}, - {'params': no_decay, 'weight_decay': 0.}] + return [{"params": has_decay}, {"params": no_decay, "weight_decay": 0.0}] def get_swin_layer(name, num_layers, depths): @@ -107,27 +150,33 @@ def get_swin_layer(name, num_layers, depths): elif name.startswith("patch_embed"): return 0 elif name.startswith("layers"): - layer_id = int(name.split('.')[1]) - block_id = name.split('.')[3] - if block_id == 'reduction' or block_id == 'norm': - return sum(depths[:layer_id + 1]) + layer_id = int(name.split(".")[1]) + block_id = name.split(".")[3] + if block_id == "reduction" or block_id == "norm": + return sum(depths[: layer_id + 1]) layer_id = sum(depths[:layer_id]) + int(block_id) return layer_id + 1 else: return num_layers - 1 -def get_finetune_param_groups(model, lr, weight_decay, get_layer_func, scales, skip_list=(), skip_keywords=()): +def get_finetune_param_groups( + model, lr, weight_decay, get_layer_func, scales, skip_list=(), skip_keywords=() +): parameter_group_names = {} parameter_group_vars = {} for name, param in model.named_parameters(): if not param.requires_grad: continue - if len(param.shape) == 1 or name.endswith(".bias") or (name in skip_list) or \ - check_keywords_in_name(name, skip_keywords): + if ( + len(param.shape) == 1 + or name.endswith(".bias") + or (name in skip_list) + or check_keywords_in_name(name, skip_keywords) + ): group_name = "no_decay" - this_weight_decay = 0. + this_weight_decay = 0.0 else: group_name = "decay" this_weight_decay = weight_decay @@ -141,7 +190,7 @@ def get_finetune_param_groups(model, lr, weight_decay, get_layer_func, scales, s if scales is not None: scale = scales[layer_id] else: - scale = 1. + scale = 1.0 parameter_group_names[group_name] = { "group_name": group_name, @@ -155,7 +204,7 @@ def get_finetune_param_groups(model, lr, weight_decay, get_layer_func, scales, s "weight_decay": this_weight_decay, "params": [], "lr": lr * scale, - "lr_scale": scale + "lr_scale": scale, } parameter_group_vars[group_name]["params"].append(param) diff --git a/utils.py b/utils.py index eb607cfe..8f09bb35 100644 --- a/utils.py +++ b/utils.py @@ -6,32 +6,43 @@ # -------------------------------------------------------- import os + import torch import torch.distributed as dist from torch._six import inf def load_checkpoint(config, model, optimizer, lr_scheduler, loss_scaler, logger): - logger.info(f"==============> Resuming form {config.MODEL.RESUME}....................") - if config.MODEL.RESUME.startswith('https'): + logger.info( + f"==============> Resuming form {config.MODEL.RESUME}...................." + ) + if config.MODEL.RESUME.startswith("https"): checkpoint = torch.hub.load_state_dict_from_url( - config.MODEL.RESUME, map_location='cpu', check_hash=True) + config.MODEL.RESUME, map_location="cpu", check_hash=True + ) else: - checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu') - msg = model.load_state_dict(checkpoint['model'], strict=False) + checkpoint = torch.load(config.MODEL.RESUME, map_location="cpu") + msg = model.load_state_dict(checkpoint["model"], strict=False) logger.info(msg) max_accuracy = 0.0 - if not config.EVAL_MODE and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint: - optimizer.load_state_dict(checkpoint['optimizer']) - lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) + if ( + not config.EVAL_MODE + and "optimizer" in checkpoint + and "lr_scheduler" in checkpoint + and "epoch" in checkpoint + ): + optimizer.load_state_dict(checkpoint["optimizer"]) + lr_scheduler.load_state_dict(checkpoint["lr_scheduler"]) config.defrost() - config.TRAIN.START_EPOCH = checkpoint['epoch'] + 1 + config.TRAIN.START_EPOCH = checkpoint["epoch"] + 1 config.freeze() - if 'scaler' in checkpoint: - loss_scaler.load_state_dict(checkpoint['scaler']) - logger.info(f"=> loaded successfully '{config.MODEL.RESUME}' (epoch {checkpoint['epoch']})") - if 'max_accuracy' in checkpoint: - max_accuracy = checkpoint['max_accuracy'] + if "scaler" in checkpoint: + loss_scaler.load_state_dict(checkpoint["scaler"]) + logger.info( + f"=> loaded successfully '{config.MODEL.RESUME}' (epoch {checkpoint['epoch']})" + ) + if "max_accuracy" in checkpoint: + max_accuracy = checkpoint["max_accuracy"] del checkpoint torch.cuda.empty_cache() @@ -39,17 +50,23 @@ def load_checkpoint(config, model, optimizer, lr_scheduler, loss_scaler, logger) def load_pretrained(config, model, logger): - logger.info(f"==============> Loading weight {config.MODEL.PRETRAINED} for fine-tuning......") - checkpoint = torch.load(config.MODEL.PRETRAINED, map_location='cpu') - state_dict = checkpoint['model'] + logger.info( + f"==============> Loading weight {config.MODEL.PRETRAINED} for fine-tuning......" + ) + checkpoint = torch.load(config.MODEL.PRETRAINED, map_location="cpu") + state_dict = checkpoint["model"] # delete relative_position_index since we always re-init it - relative_position_index_keys = [k for k in state_dict.keys() if "relative_position_index" in k] + relative_position_index_keys = [ + k for k in state_dict.keys() if "relative_position_index" in k + ] for k in relative_position_index_keys: del state_dict[k] # delete relative_coords_table since we always re-init it - relative_position_index_keys = [k for k in state_dict.keys() if "relative_coords_table" in k] + relative_position_index_keys = [ + k for k in state_dict.keys() if "relative_coords_table" in k + ] for k in relative_position_index_keys: del state_dict[k] @@ -59,7 +76,9 @@ def load_pretrained(config, model, logger): del state_dict[k] # bicubic interpolate relative_position_bias_table if not match - relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k] + relative_position_bias_table_keys = [ + k for k in state_dict.keys() if "relative_position_bias_table" in k + ] for k in relative_position_bias_table_keys: relative_position_bias_table_pretrained = state_dict[k] relative_position_bias_table_current = model.state_dict()[k] @@ -70,15 +89,25 @@ def load_pretrained(config, model, logger): else: if L1 != L2: # bicubic interpolate relative_position_bias_table if not match - S1 = int(L1 ** 0.5) - S2 = int(L2 ** 0.5) - relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate( - relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1), size=(S2, S2), - mode='bicubic') - state_dict[k] = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0) + S1 = int(L1**0.5) + S2 = int(L2**0.5) + relative_position_bias_table_pretrained_resized = ( + torch.nn.functional.interpolate( + relative_position_bias_table_pretrained.permute(1, 0).view( + 1, nH1, S1, S1 + ), + size=(S2, S2), + mode="bicubic", + ) + ) + state_dict[k] = relative_position_bias_table_pretrained_resized.view( + nH2, L2 + ).permute(1, 0) # bicubic interpolate absolute_pos_embed if not match - absolute_pos_embed_keys = [k for k in state_dict.keys() if "absolute_pos_embed" in k] + absolute_pos_embed_keys = [ + k for k in state_dict.keys() if "absolute_pos_embed" in k + ] for k in absolute_pos_embed_keys: # dpe absolute_pos_embed_pretrained = state_dict[k] @@ -89,35 +118,46 @@ def load_pretrained(config, model, logger): logger.warning(f"Error in loading {k}, passing......") else: if L1 != L2: - S1 = int(L1 ** 0.5) - S2 = int(L2 ** 0.5) - absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape(-1, S1, S1, C1) - absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute(0, 3, 1, 2) + S1 = int(L1**0.5) + S2 = int(L2**0.5) + absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape( + -1, S1, S1, C1 + ) + absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute( + 0, 3, 1, 2 + ) absolute_pos_embed_pretrained_resized = torch.nn.functional.interpolate( - absolute_pos_embed_pretrained, size=(S2, S2), mode='bicubic') - absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1) - absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.flatten(1, 2) + absolute_pos_embed_pretrained, size=(S2, S2), mode="bicubic" + ) + absolute_pos_embed_pretrained_resized = ( + absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1) + ) + absolute_pos_embed_pretrained_resized = ( + absolute_pos_embed_pretrained_resized.flatten(1, 2) + ) state_dict[k] = absolute_pos_embed_pretrained_resized # check classifier, if not match, then re-init classifier to zero - head_bias_pretrained = state_dict['head.bias'] + head_bias_pretrained = state_dict["head.bias"] Nc1 = head_bias_pretrained.shape[0] Nc2 = model.head.bias.shape[0] - if (Nc1 != Nc2): + if Nc1 != Nc2: if Nc1 == 21841 and Nc2 == 1000: logger.info("loading ImageNet-22K weight to ImageNet-1K ......") - map22kto1k_path = f'data/map22kto1k.txt' + map22kto1k_path = f"data/map22kto1k.txt" with open(map22kto1k_path) as f: map22kto1k = f.readlines() map22kto1k = [int(id22k.strip()) for id22k in map22kto1k] - state_dict['head.weight'] = state_dict['head.weight'][map22kto1k, :] - state_dict['head.bias'] = state_dict['head.bias'][map22kto1k] + state_dict["head.weight"] = state_dict["head.weight"][map22kto1k, :] + state_dict["head.bias"] = state_dict["head.bias"][map22kto1k] else: - torch.nn.init.constant_(model.head.bias, 0.) - torch.nn.init.constant_(model.head.weight, 0.) - del state_dict['head.weight'] - del state_dict['head.bias'] - logger.warning(f"Error in loading classifier head, re-init classifier head to 0") + torch.nn.init.constant_(model.head.bias, 0.0) + torch.nn.init.constant_(model.head.weight, 0.0) + del state_dict["head.weight"] + del state_dict["head.bias"] + logger.warning( + f"Error in loading classifier head, re-init classifier head to 0" + ) msg = model.load_state_dict(state_dict, strict=False) logger.warning(msg) @@ -128,16 +168,20 @@ def load_pretrained(config, model, logger): torch.cuda.empty_cache() -def save_checkpoint(config, epoch, model, max_accuracy, optimizer, lr_scheduler, loss_scaler, logger): - save_state = {'model': model.state_dict(), - 'optimizer': optimizer.state_dict(), - 'lr_scheduler': lr_scheduler.state_dict(), - 'max_accuracy': max_accuracy, - 'scaler': loss_scaler.state_dict(), - 'epoch': epoch, - 'config': config} - - save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth') +def save_checkpoint( + config, epoch, model, max_accuracy, optimizer, lr_scheduler, loss_scaler, logger +): + save_state = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "lr_scheduler": lr_scheduler.state_dict(), + "max_accuracy": max_accuracy, + "scaler": loss_scaler.state_dict(), + "epoch": epoch, + "config": config, + } + + save_path = os.path.join(config.OUTPUT, f"ckpt_epoch_{epoch}.pth") logger.info(f"{save_path} saving......") torch.save(save_state, save_path) logger.info(f"{save_path} saved !!!") @@ -152,16 +196,18 @@ def get_grad_norm(parameters, norm_type=2): for p in parameters: param_norm = p.grad.data.norm(norm_type) total_norm += param_norm.item() ** norm_type - total_norm = total_norm ** (1. / norm_type) + total_norm = total_norm ** (1.0 / norm_type) return total_norm def auto_resume_helper(output_dir): checkpoints = os.listdir(output_dir) - checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith('pth')] + checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith("pth")] print(f"All checkpoints founded in {output_dir}: {checkpoints}") if len(checkpoints) > 0: - latest_checkpoint = max([os.path.join(output_dir, d) for d in checkpoints], key=os.path.getmtime) + latest_checkpoint = max( + [os.path.join(output_dir, d) for d in checkpoints], key=os.path.getmtime + ) print(f"The latest checkpoint founded: {latest_checkpoint}") resume_file = latest_checkpoint else: @@ -182,13 +228,17 @@ def ampscaler_get_grad_norm(parameters, norm_type: float = 2.0) -> torch.Tensor: parameters = [p for p in parameters if p.grad is not None] norm_type = float(norm_type) if len(parameters) == 0: - return torch.tensor(0.) + return torch.tensor(0.0) device = parameters[0].grad.device if norm_type == inf: total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters) else: - total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), - norm_type).to(device) for p in parameters]), norm_type) + total_norm = torch.norm( + torch.stack( + [torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters] + ), + norm_type, + ) return total_norm @@ -198,12 +248,22 @@ class NativeScalerWithGradNormCount: def __init__(self): self._scaler = torch.cuda.amp.GradScaler() - def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True): + def __call__( + self, + loss, + optimizer, + clip_grad=None, + parameters=None, + create_graph=False, + update_grad=True, + ): self._scaler.scale(loss).backward(create_graph=create_graph) if update_grad: if clip_grad is not None: assert parameters is not None - self._scaler.unscale_(optimizer) # unscale the gradients of optimizer's assigned params in-place + self._scaler.unscale_( + optimizer + ) # unscale the gradients of optimizer's assigned params in-place norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad) else: self._scaler.unscale_(optimizer) diff --git a/utils_moe.py b/utils_moe.py index c5368a1a..c4c9f3a7 100644 --- a/utils_moe.py +++ b/utils_moe.py @@ -6,6 +6,7 @@ # -------------------------------------------------------- import os + import torch import torch.distributed as dist @@ -30,31 +31,42 @@ def merge_moe_model_state_dict(moe_model_state_dict, non_moe_model_state_dict): def load_checkpoint(config, model, optimizer, lr_scheduler, loss_scaler, logger): global_rank = dist.get_rank() - logger.info(f"==============> Rank[{global_rank}] Resuming form {config.MODEL.RESUME}....................") - if config.MODEL.RESUME.endswith(f'.pth'): + logger.info( + f"==============> Rank[{global_rank}] Resuming form {config.MODEL.RESUME}...................." + ) + if config.MODEL.RESUME.endswith(f".pth"): if config.TRAIN.MOE.SAVE_MASTER: - resume_path = config.MODEL.RESUME + f'.global' + resume_path = config.MODEL.RESUME + f".global" else: - resume_path = config.MODEL.RESUME + f'.rank{global_rank}' - logger.info(f"===> Rank[{global_rank}] Re-formatting checkpoint name to {resume_path}......") + resume_path = config.MODEL.RESUME + f".rank{global_rank}" + logger.info( + f"===> Rank[{global_rank}] Re-formatting checkpoint name to {resume_path}......" + ) else: resume_path = config.MODEL.RESUME - checkpoint = torch.load(resume_path, map_location='cpu') - msg = model.load_state_dict(checkpoint['model'], strict=False) + checkpoint = torch.load(resume_path, map_location="cpu") + msg = model.load_state_dict(checkpoint["model"], strict=False) logger.info(msg) max_accuracy = 0.0 - if not config.EVAL_MODE and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint: - optimizer.load_state_dict(checkpoint['optimizer']) - lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) + if ( + not config.EVAL_MODE + and "optimizer" in checkpoint + and "lr_scheduler" in checkpoint + and "epoch" in checkpoint + ): + optimizer.load_state_dict(checkpoint["optimizer"]) + lr_scheduler.load_state_dict(checkpoint["lr_scheduler"]) config.defrost() - config.TRAIN.START_EPOCH = checkpoint['epoch'] + 1 + config.TRAIN.START_EPOCH = checkpoint["epoch"] + 1 config.freeze() - if 'scaler' in checkpoint: - loss_scaler.load_state_dict(checkpoint['scaler']) - logger.info(f"=>Rank[{global_rank}] loaded successfully '{config.MODEL.RESUME}' (epoch {checkpoint['epoch']})") - if 'max_accuracy' in checkpoint: - max_accuracy = checkpoint['max_accuracy'] + if "scaler" in checkpoint: + loss_scaler.load_state_dict(checkpoint["scaler"]) + logger.info( + f"=>Rank[{global_rank}] loaded successfully '{config.MODEL.RESUME}' (epoch {checkpoint['epoch']})" + ) + if "max_accuracy" in checkpoint: + max_accuracy = checkpoint["max_accuracy"] del checkpoint torch.cuda.empty_cache() @@ -63,37 +75,49 @@ def load_checkpoint(config, model, optimizer, lr_scheduler, loss_scaler, logger) def load_pretrained(config, model, logger): global_rank = dist.get_rank() - logger.info(f"==============> Rank[{global_rank}] Loading weight {config.MODEL.PRETRAINED} for fine-tuning......") - if config.MODEL.PRETRAINED.endswith(f'.pth'): + logger.info( + f"==============> Rank[{global_rank}] Loading weight {config.MODEL.PRETRAINED} for fine-tuning......" + ) + if config.MODEL.PRETRAINED.endswith(f".pth"): if config.TRAIN.MOE.SAVE_MASTER: - pretrained_path = config.MODEL.PRETRAINED + f'.global' + pretrained_path = config.MODEL.PRETRAINED + f".global" else: - pretrained_path = config.MODEL.PRETRAINED + f'.rank{global_rank}' - logger.info(f"===> Rank[{global_rank}] Re-formatting checkpoint name to {pretrained_path}......") + pretrained_path = config.MODEL.PRETRAINED + f".rank{global_rank}" + logger.info( + f"===> Rank[{global_rank}] Re-formatting checkpoint name to {pretrained_path}......" + ) else: pretrained_path = config.MODEL.PRETRAINED - if pretrained_path.endswith(f'.rank{global_rank}'): - checkpoint = torch.load(pretrained_path, map_location='cpu') - if os.path.exists(pretrained_path.replace(f'.rank{global_rank}', f'.master')): - checkpoint_master = torch.load(pretrained_path.replace(f'.rank{global_rank}', f'.master'), - map_location='cpu') - state_dict = merge_moe_model_state_dict(checkpoint['model'], checkpoint_master['model']) + if pretrained_path.endswith(f".rank{global_rank}"): + checkpoint = torch.load(pretrained_path, map_location="cpu") + if os.path.exists(pretrained_path.replace(f".rank{global_rank}", f".master")): + checkpoint_master = torch.load( + pretrained_path.replace(f".rank{global_rank}", f".master"), + map_location="cpu", + ) + state_dict = merge_moe_model_state_dict( + checkpoint["model"], checkpoint_master["model"] + ) else: - state_dict = checkpoint['model'] - elif pretrained_path.endswith(f'.pth.global'): - checkpoint = torch.load(pretrained_path, map_location='cpu') - state_dict = checkpoint['model'] + state_dict = checkpoint["model"] + elif pretrained_path.endswith(f".pth.global"): + checkpoint = torch.load(pretrained_path, map_location="cpu") + state_dict = checkpoint["model"] else: raise NotImplementedError(f"{config.MODEL.PRETRAINED} file error...") # delete relative_position_index since we always re-init it - relative_position_index_keys = [k for k in state_dict.keys() if "relative_position_index" in k] + relative_position_index_keys = [ + k for k in state_dict.keys() if "relative_position_index" in k + ] for k in relative_position_index_keys: del state_dict[k] # delete relative_coords_table since we always re-init it - relative_position_index_keys = [k for k in state_dict.keys() if "relative_coords_table" in k] + relative_position_index_keys = [ + k for k in state_dict.keys() if "relative_coords_table" in k + ] for k in relative_position_index_keys: del state_dict[k] @@ -103,7 +127,9 @@ def load_pretrained(config, model, logger): del state_dict[k] # bicubic interpolate relative_position_bias_table if not match - relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k] + relative_position_bias_table_keys = [ + k for k in state_dict.keys() if "relative_position_bias_table" in k + ] for k in relative_position_bias_table_keys: relative_position_bias_table_pretrained = state_dict[k] relative_position_bias_table_current = model.state_dict()[k] @@ -114,15 +140,25 @@ def load_pretrained(config, model, logger): else: if L1 != L2: # bicubic interpolate relative_position_bias_table if not match - S1 = int(L1 ** 0.5) - S2 = int(L2 ** 0.5) - relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate( - relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1), size=(S2, S2), - mode='bicubic') - state_dict[k] = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0) + S1 = int(L1**0.5) + S2 = int(L2**0.5) + relative_position_bias_table_pretrained_resized = ( + torch.nn.functional.interpolate( + relative_position_bias_table_pretrained.permute(1, 0).view( + 1, nH1, S1, S1 + ), + size=(S2, S2), + mode="bicubic", + ) + ) + state_dict[k] = relative_position_bias_table_pretrained_resized.view( + nH2, L2 + ).permute(1, 0) # bicubic interpolate absolute_pos_embed if not match - absolute_pos_embed_keys = [k for k in state_dict.keys() if "absolute_pos_embed" in k] + absolute_pos_embed_keys = [ + k for k in state_dict.keys() if "absolute_pos_embed" in k + ] for k in absolute_pos_embed_keys: # dpe absolute_pos_embed_pretrained = state_dict[k] @@ -133,35 +169,46 @@ def load_pretrained(config, model, logger): logger.warning(f"Error in loading {k}, passing......") else: if L1 != L2: - S1 = int(L1 ** 0.5) - S2 = int(L2 ** 0.5) - absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape(-1, S1, S1, C1) - absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute(0, 3, 1, 2) + S1 = int(L1**0.5) + S2 = int(L2**0.5) + absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape( + -1, S1, S1, C1 + ) + absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute( + 0, 3, 1, 2 + ) absolute_pos_embed_pretrained_resized = torch.nn.functional.interpolate( - absolute_pos_embed_pretrained, size=(S2, S2), mode='bicubic') - absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1) - absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.flatten(1, 2) + absolute_pos_embed_pretrained, size=(S2, S2), mode="bicubic" + ) + absolute_pos_embed_pretrained_resized = ( + absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1) + ) + absolute_pos_embed_pretrained_resized = ( + absolute_pos_embed_pretrained_resized.flatten(1, 2) + ) state_dict[k] = absolute_pos_embed_pretrained_resized # check classifier, if not match, then re-init classifier to zero - head_bias_pretrained = state_dict['head.bias'] + head_bias_pretrained = state_dict["head.bias"] Nc1 = head_bias_pretrained.shape[0] Nc2 = model.head.bias.shape[0] - if (Nc1 != Nc2): + if Nc1 != Nc2: if Nc1 == 21841 and Nc2 == 1000: logger.info("loading ImageNet-22K weight to ImageNet-1K ......") - map22kto1k_path = f'data/map22kto1k.txt' + map22kto1k_path = f"data/map22kto1k.txt" with open(map22kto1k_path) as f: map22kto1k = f.readlines() map22kto1k = [int(id22k.strip()) for id22k in map22kto1k] - state_dict['head.weight'] = state_dict['head.weight'][map22kto1k, :] - state_dict['head.bias'] = state_dict['head.bias'][map22kto1k] + state_dict["head.weight"] = state_dict["head.weight"][map22kto1k, :] + state_dict["head.bias"] = state_dict["head.bias"][map22kto1k] else: - torch.nn.init.constant_(model.head.bias, 0.) - torch.nn.init.constant_(model.head.weight, 0.) - del state_dict['head.weight'] - del state_dict['head.bias'] - logger.warning(f"Error in loading classifier head, re-init classifier head to 0") + torch.nn.init.constant_(model.head.bias, 0.0) + torch.nn.init.constant_(model.head.weight, 0.0) + del state_dict["head.weight"] + del state_dict["head.bias"] + logger.warning( + f"Error in loading classifier head, re-init classifier head to 0" + ) msg = model.load_state_dict(state_dict, strict=False) logger.warning(msg) @@ -172,48 +219,70 @@ def load_pretrained(config, model, logger): torch.cuda.empty_cache() -def save_checkpoint(config, epoch, model, max_accuracy, optimizer, lr_scheduler, loss_scaler, logger, - zero_redundancy=False): +def save_checkpoint( + config, + epoch, + model, + max_accuracy, + optimizer, + lr_scheduler, + loss_scaler, + logger, + zero_redundancy=False, +): global_rank = dist.get_rank() if zero_redundancy: if config.TRAIN.MOE.SAVE_MASTER: - save_state = {'model': model.state_dict()} + save_state = {"model": model.state_dict()} if global_rank == 0: - save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth.global') + save_path = os.path.join( + config.OUTPUT, f"ckpt_epoch_{epoch}.pth.global" + ) logger.info(f"{save_path} saving......") torch.save(save_state, save_path) logger.info(f"{save_path} saved !!!") else: - moe_model_state_dict, non_moe_model_state_dict = \ - split_moe_model_state_dict(model._ddp_params_and_buffers_to_ignore, model.state_dict()) - save_state = {'model': moe_model_state_dict} - save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth.rank{global_rank}') + moe_model_state_dict, non_moe_model_state_dict = split_moe_model_state_dict( + model._ddp_params_and_buffers_to_ignore, model.state_dict() + ) + save_state = {"model": moe_model_state_dict} + save_path = os.path.join( + config.OUTPUT, f"ckpt_epoch_{epoch}.pth.rank{global_rank}" + ) logger.info(f"{save_path} saving......") torch.save(save_state, save_path) logger.info(f"{save_path} saved !!!") if global_rank == 0: - save_state_master = {'model': non_moe_model_state_dict} - save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth.master') + save_state_master = {"model": non_moe_model_state_dict} + save_path = os.path.join( + config.OUTPUT, f"ckpt_epoch_{epoch}.pth.master" + ) logger.info(f"{save_path} saving......") torch.save(save_state_master, save_path) logger.info(f"{save_path} saved !!!") else: - save_state = {'model': model.state_dict(), - 'optimizer': optimizer.state_dict(), - 'lr_scheduler': lr_scheduler.state_dict(), - 'max_accuracy': max_accuracy, - 'scaler': loss_scaler.state_dict(), - 'epoch': epoch, - 'config': config} + save_state = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "lr_scheduler": lr_scheduler.state_dict(), + "max_accuracy": max_accuracy, + "scaler": loss_scaler.state_dict(), + "epoch": epoch, + "config": config, + } if config.TRAIN.MOE.SAVE_MASTER: if global_rank == 0: - save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth.global') + save_path = os.path.join( + config.OUTPUT, f"ckpt_epoch_{epoch}.pth.global" + ) logger.info(f"{save_path} saving......") torch.save(save_state, save_path) logger.info(f"{save_path} saved !!!") else: - save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth.rank{global_rank}') + save_path = os.path.join( + config.OUTPUT, f"ckpt_epoch_{epoch}.pth.rank{global_rank}" + ) logger.info(f"{save_path} saving......") torch.save(save_state, save_path) logger.info(f"{save_path} saved !!!") @@ -223,13 +292,22 @@ def auto_resume_helper(output_dir, save_master=False): global_rank = dist.get_rank() checkpoints = os.listdir(output_dir) if not save_master: - master_checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith(f'pth.rank0')] + master_checkpoints = [ + ckpt for ckpt in checkpoints if ckpt.endswith(f"pth.rank0") + ] else: - master_checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith(f'pth.global')] + master_checkpoints = [ + ckpt for ckpt in checkpoints if ckpt.endswith(f"pth.global") + ] print(f"All master checkpoints founded in {output_dir}: {master_checkpoints}") if len(master_checkpoints) > 0: - latest_master_checkpoint = max([os.path.join(output_dir, d) for d in master_checkpoints], key=os.path.getmtime) - latest_checkpoint = latest_master_checkpoint.replace('pth.rank0', f'pth.rank{global_rank}') + latest_master_checkpoint = max( + [os.path.join(output_dir, d) for d in master_checkpoints], + key=os.path.getmtime, + ) + latest_checkpoint = latest_master_checkpoint.replace( + "pth.rank0", f"pth.rank{global_rank}" + ) print(f"The latest checkpoint founded: {latest_checkpoint}") resume_file = latest_checkpoint else: diff --git a/utils_simmim.py b/utils_simmim.py index 2643fa44..2b7b7ec2 100644 --- a/utils_simmim.py +++ b/utils_simmim.py @@ -7,41 +7,53 @@ # -------------------------------------------------------- import os + +import numpy as np import torch import torch.distributed as dist -import numpy as np from scipy import interpolate def load_checkpoint(config, model, optimizer, lr_scheduler, scaler, logger): logger.info(f">>>>>>>>>> Resuming from {config.MODEL.RESUME} ..........") - if config.MODEL.RESUME.startswith('https'): + if config.MODEL.RESUME.startswith("https"): checkpoint = torch.hub.load_state_dict_from_url( - config.MODEL.RESUME, map_location='cpu', check_hash=True) + config.MODEL.RESUME, map_location="cpu", check_hash=True + ) else: - checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu') + checkpoint = torch.load(config.MODEL.RESUME, map_location="cpu") # re-map keys due to name change (only for loading provided models) - rpe_mlp_keys = [k for k in checkpoint['model'].keys() if "rpe_mlp" in k] + rpe_mlp_keys = [k for k in checkpoint["model"].keys() if "rpe_mlp" in k] for k in rpe_mlp_keys: - checkpoint['model'][k.replace('rpe_mlp', 'cpb_mlp')] = checkpoint['model'].pop(k) - - msg = model.load_state_dict(checkpoint['model'], strict=False) + checkpoint["model"][k.replace("rpe_mlp", "cpb_mlp")] = checkpoint["model"].pop( + k + ) + + msg = model.load_state_dict(checkpoint["model"], strict=False) logger.info(msg) max_accuracy = 0.0 - if not config.EVAL_MODE and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'scaler' in checkpoint and 'epoch' in checkpoint: - optimizer.load_state_dict(checkpoint['optimizer']) - lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) - scaler.load_state_dict(checkpoint['scaler']) + if ( + not config.EVAL_MODE + and "optimizer" in checkpoint + and "lr_scheduler" in checkpoint + and "scaler" in checkpoint + and "epoch" in checkpoint + ): + optimizer.load_state_dict(checkpoint["optimizer"]) + lr_scheduler.load_state_dict(checkpoint["lr_scheduler"]) + scaler.load_state_dict(checkpoint["scaler"]) config.defrost() - config.TRAIN.START_EPOCH = checkpoint['epoch'] + 1 + config.TRAIN.START_EPOCH = checkpoint["epoch"] + 1 config.freeze() - logger.info(f"=> loaded successfully '{config.MODEL.RESUME}' (epoch {checkpoint['epoch']})") - if 'max_accuracy' in checkpoint: - max_accuracy = checkpoint['max_accuracy'] + logger.info( + f"=> loaded successfully '{config.MODEL.RESUME}' (epoch {checkpoint['epoch']})" + ) + if "max_accuracy" in checkpoint: + max_accuracy = checkpoint["max_accuracy"] else: max_accuracy = 0.0 @@ -50,16 +62,20 @@ def load_checkpoint(config, model, optimizer, lr_scheduler, scaler, logger): return max_accuracy -def save_checkpoint(config, epoch, model, max_accuracy, optimizer, lr_scheduler, scaler, logger): - save_state = {'model': model.state_dict(), - 'optimizer': optimizer.state_dict(), - 'lr_scheduler': lr_scheduler.state_dict(), - 'scaler': scaler.state_dict(), - 'max_accuracy': max_accuracy, - 'epoch': epoch, - 'config': config} - - save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth') +def save_checkpoint( + config, epoch, model, max_accuracy, optimizer, lr_scheduler, scaler, logger +): + save_state = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "lr_scheduler": lr_scheduler.state_dict(), + "scaler": scaler.state_dict(), + "max_accuracy": max_accuracy, + "epoch": epoch, + "config": config, + } + + save_path = os.path.join(config.OUTPUT, f"ckpt_epoch_{epoch}.pth") logger.info(f"{save_path} saving......") torch.save(save_state, save_path) logger.info(f"{save_path} saved !!!") @@ -74,16 +90,18 @@ def get_grad_norm(parameters, norm_type=2): for p in parameters: param_norm = p.grad.data.norm(norm_type) total_norm += param_norm.item() ** norm_type - total_norm = total_norm ** (1. / norm_type) + total_norm = total_norm ** (1.0 / norm_type) return total_norm def auto_resume_helper(output_dir, logger): checkpoints = os.listdir(output_dir) - checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith('pth')] + checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith("pth")] logger.info(f"All checkpoints founded in {output_dir}: {checkpoints}") if len(checkpoints) > 0: - latest_checkpoint = max([os.path.join(output_dir, d) for d in checkpoints], key=os.path.getmtime) + latest_checkpoint = max( + [os.path.join(output_dir, d) for d in checkpoints], key=os.path.getmtime + ) logger.info(f"The latest checkpoint founded: {latest_checkpoint}") resume_file = latest_checkpoint else: @@ -100,16 +118,20 @@ def reduce_tensor(tensor): def load_pretrained(config, model, logger): logger.info(f">>>>>>>>>> Fine-tuned from {config.MODEL.PRETRAINED} ..........") - checkpoint = torch.load(config.MODEL.PRETRAINED, map_location='cpu') - checkpoint_model = checkpoint['model'] - - if any([True if 'encoder.' in k else False for k in checkpoint_model.keys()]): - checkpoint_model = {k.replace('encoder.', ''): v for k, v in checkpoint_model.items() if k.startswith('encoder.')} - logger.info('Detect pre-trained model, remove [encoder.] prefix.') + checkpoint = torch.load(config.MODEL.PRETRAINED, map_location="cpu") + checkpoint_model = checkpoint["model"] + + if any([True if "encoder." in k else False for k in checkpoint_model.keys()]): + checkpoint_model = { + k.replace("encoder.", ""): v + for k, v in checkpoint_model.items() + if k.startswith("encoder.") + } + logger.info("Detect pre-trained model, remove [encoder.] prefix.") else: - logger.info('Detect non-pre-trained model, pass without doing anything.') + logger.info("Detect non-pre-trained model, pass without doing anything.") - if config.MODEL.TYPE in ['swin', 'swinv2']: + if config.MODEL.TYPE in ["swin", "swinv2"]: logger.info(f">>>>>>>>>> Remapping pre-trained keys for SWIN ..........") checkpoint = remap_pretrained_keys_swin(model, checkpoint_model, logger) else: @@ -117,15 +139,15 @@ def load_pretrained(config, model, logger): msg = model.load_state_dict(checkpoint_model, strict=False) logger.info(msg) - + del checkpoint torch.cuda.empty_cache() logger.info(f">>>>>>>>>> loaded successfully '{config.MODEL.PRETRAINED}'") - + def remap_pretrained_keys_swin(model, checkpoint_model, logger): state_dict = model.state_dict() - + # Geometric interpolation when pre-trained patch size mismatch with fine-tuned patch size all_keys = list(checkpoint_model.keys()) for key in all_keys: @@ -138,12 +160,14 @@ def remap_pretrained_keys_swin(model, checkpoint_model, logger): logger.info(f"Error in loading {key}, passing......") else: if L1 != L2: - logger.info(f"{key}: Interpolate relative_position_bias_table using geo.") - src_size = int(L1 ** 0.5) - dst_size = int(L2 ** 0.5) + logger.info( + f"{key}: Interpolate relative_position_bias_table using geo." + ) + src_size = int(L1**0.5) + dst_size = int(L2**0.5) def geometric_progression(a, r, n): - return a * (1.0 - r ** n) / (1.0 - r) + return a * (1.0 - r**n) / (1.0 - r) left, right = 1.01, 1.5 while right - left > 1e-6: @@ -178,28 +202,41 @@ def geometric_progression(a, r, n): all_rel_pos_bias = [] for i in range(nH1): - z = relative_position_bias_table_pretrained[:, i].view(src_size, src_size).float().numpy() - f_cubic = interpolate.interp2d(x, y, z, kind='cubic') - all_rel_pos_bias.append(torch.Tensor(f_cubic(dx, dy)).contiguous().view(-1, 1).to( - relative_position_bias_table_pretrained.device)) + z = ( + relative_position_bias_table_pretrained[:, i] + .view(src_size, src_size) + .float() + .numpy() + ) + f_cubic = interpolate.interp2d(x, y, z, kind="cubic") + all_rel_pos_bias.append( + torch.Tensor(f_cubic(dx, dy)) + .contiguous() + .view(-1, 1) + .to(relative_position_bias_table_pretrained.device) + ) new_rel_pos_bias = torch.cat(all_rel_pos_bias, dim=-1) checkpoint_model[key] = new_rel_pos_bias # delete relative_position_index since we always re-init it - relative_position_index_keys = [k for k in checkpoint_model.keys() if "relative_position_index" in k] + relative_position_index_keys = [ + k for k in checkpoint_model.keys() if "relative_position_index" in k + ] for k in relative_position_index_keys: del checkpoint_model[k] # delete relative_coords_table since we always re-init it - relative_coords_table_keys = [k for k in checkpoint_model.keys() if "relative_coords_table" in k] + relative_coords_table_keys = [ + k for k in checkpoint_model.keys() if "relative_coords_table" in k + ] for k in relative_coords_table_keys: del checkpoint_model[k] # re-map keys due to name change rpe_mlp_keys = [k for k in checkpoint_model.keys() if "rpe_mlp" in k] for k in rpe_mlp_keys: - checkpoint_model[k.replace('rpe_mlp', 'cpb_mlp')] = checkpoint_model.pop(k) + checkpoint_model[k.replace("rpe_mlp", "cpb_mlp")] = checkpoint_model.pop(k) # delete attn_mask since we always re-init it attn_mask_keys = [k for k in checkpoint_model.keys() if "attn_mask" in k] From 477d019c07a1af45e798ecfeabae273692649323 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Wed, 19 Oct 2022 18:08:05 +0000 Subject: [PATCH 02/47] Add software engineering tools --- makefile | 6 ++++++ setup.cfg | 5 +++++ 2 files changed, 11 insertions(+) create mode 100644 makefile create mode 100644 setup.cfg diff --git a/makefile b/makefile new file mode 100644 index 00000000..caf42d01 --- /dev/null +++ b/makefile @@ -0,0 +1,6 @@ +fmt: + fd -e py | xargs isort --profile black + fd -e py | xargs black + +lint: + fd -e py | xargs flake8 diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 00000000..4d53e908 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,5 @@ +[flake8] +ignore = E501,E203,E722,W503,W391 + +[pycodestyle] +ignore = E501,E203,E722,W503,W391 From e0ddde84407664e37a6ec6ec7eced2658b2b30b5 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Tue, 25 Oct 2022 01:33:48 +0000 Subject: [PATCH 03/47] Added hierarchical stuff --- .gitignore | 5 ++ config.py | 19 +++++- data/build.py | 43 ++++++++++++- data/constants.py | 8 +++ data/hierarchical.py | 94 ++++++++++++++++++++++++++++ hierarchical.py | 77 +++++++++++++++++++++++ logger.py | 28 +++++++++ lr_scheduler.py | 77 +++++++++++++++-------- main.py | 111 +++++++++++++++++++++++++--------- models/build.py | 2 +- models/swin_transformer_v2.py | 53 +++++++++++++--- optimizer.py | 2 +- scripts/train.fish | 56 +++++++++++++++++ utils.py | 11 +++- 14 files changed, 514 insertions(+), 72 deletions(-) create mode 100644 data/constants.py create mode 100644 data/hierarchical.py create mode 100644 hierarchical.py create mode 100644 scripts/train.fish diff --git a/.gitignore b/.gitignore index 24f67f2e..8e3e1076 100644 --- a/.gitignore +++ b/.gitignore @@ -133,3 +133,8 @@ dmypy.json # Pyre type checker .pyre/ + +# Apex +apex/ + +runs/ diff --git a/config.py b/config.py index e5ba8761..aafdfdf0 100644 --- a/config.py +++ b/config.py @@ -194,6 +194,13 @@ _C.TRAIN.MOE = CN() # Only save model on master device _C.TRAIN.MOE.SAVE_MASTER = False + +# Hierarchical coefficients for loss +_C.TRAIN.HIERARCHICAL_COEFFS = (1,) + +# [Debugging] How many batches of the training data to overfit. +_C.TRAIN.OVERFIT_BATCHES = 0 + # ----------------------------------------------------------------------------- # Augmentation settings # ----------------------------------------------------------------------------- @@ -234,6 +241,10 @@ # ----------------------------------------------------------------------------- # Misc # ----------------------------------------------------------------------------- + +# Whether we are doing hierarchical classification +_C.HIERARHICAL = False + # [SimMIM] Whether to enable pytorch amp, overwritten by command line argument _C.ENABLE_AMP = False @@ -330,12 +341,14 @@ def _check_args(name): config.FUSED_WINDOW_PROCESS = True if _check_args("fused_layernorm"): config.FUSED_LAYERNORM = True - ## Overwrite optimizer if not None, currently we use it for [fused_adam, fused_lamb] + # Overwrite optimizer if not None, currently we use it for [fused_adam, fused_lamb] if _check_args("optim"): config.TRAIN.OPTIMIZER.NAME = args.optim - # set local rank for distributed training - config.LOCAL_RANK = args.local_rank + # Use os.environ["LOCAL_RANK"] rather than --local_rank + if "LOCAL_RANK" in os.environ: + # set local rank for distributed training + config.LOCAL_RANK = int(os.environ["LOCAL_RANK"]) # output folder config.OUTPUT = os.path.join(config.OUTPUT, config.MODEL.NAME, config.TAG) diff --git a/data/build.py b/data/build.py index 85a93e98..6b6d3b0c 100644 --- a/data/build.py +++ b/data/build.py @@ -6,15 +6,19 @@ # -------------------------------------------------------- import os +import random import numpy as np import torch import torch.distributed as dist from timm.data import Mixup, create_transform from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD +from torch.utils.data import Subset from torchvision import datasets, transforms from .cached_image_folder import CachedImageFolder +from .constants import data_mean_std +from .hierarchical import HierarchicalImageFolder, HierarchicalMixup from .imagenet22k_dataset import IN22KDATASET from .samplers import SubsetRandomSampler @@ -35,7 +39,7 @@ def _pil_interp(method): import timm.data.transforms as timm_transforms timm_transforms._pil_interp = _pil_interp -except: +except ImportError: from timm.data.transforms import _pil_interp @@ -53,6 +57,12 @@ def build_loader(config): f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build val dataset" ) + # Check if we are overfitting some subset of the training data for debugging + if config.TRAIN.OVERFIT_BATCHES > 0: + n_examples = config.TRAIN.OVERFIT_BATCHES * config.DATA.BATCH_SIZE + indices = random.sample(range(len(dataset_train)), n_examples) + dataset_train = Subset(dataset_train, indices) + num_tasks = dist.get_world_size() global_rank = dist.get_rank() if config.DATA.ZIP_MODE and config.DATA.CACHE_MODE == "part": @@ -97,7 +107,7 @@ def build_loader(config): or config.AUG.CUTMIX_MINMAX is not None ) if mixup_active: - mixup_fn = Mixup( + mixup_args = dict( mixup_alpha=config.AUG.MIXUP, cutmix_alpha=config.AUG.CUTMIX, cutmix_minmax=config.AUG.CUTMIX_MINMAX, @@ -107,6 +117,10 @@ def build_loader(config): label_smoothing=config.MODEL.LABEL_SMOOTHING, num_classes=config.MODEL.NUM_CLASSES, ) + if config.HIERARHICAL: + mixup_fn = HierarchicalMixup(**mixup_args) + else: + mixup_fn = Mixup(**mixup_args) return dataset_train, dataset_val, data_loader_train, data_loader_val, mixup_fn @@ -137,6 +151,19 @@ def build_dataset(is_train, config): ann_file = prefix + "_map_val.txt" dataset = IN22KDATASET(config.DATA.DATA_PATH, ann_file, transform) nb_classes = 21841 + + elif config.DATA.DATASET == "inat21": + if config.DATA.ZIP_MODE: + raise NotImplementedError("We do not support zipped inat21") + + prefix = "train" if is_train else "val" + root = os.path.join(config.DATA.DATA_PATH, prefix) + if config.HIERARHICAL: + dataset = HierarchicalImageFolder(root, transform=transform) + nb_classes = dataset.num_classes + else: + dataset = datasets.ImageFolder(root, transform=transform) + nb_classes = 10_000 else: raise NotImplementedError("We only support ImageNet Now.") @@ -188,6 +215,16 @@ def build_transform(is_train, config): ) ) + if config.DATA.DATA_PATH in data_mean_std: + mean, std = data_mean_std[config.DATA.DATA_PATH] + elif config.DATA.DATASET in data_mean_std: + mean, std = data_mean_std[config.DATA.DATASET] + else: + print( + "Can't find mean/std for {config.DATA.DATASET} at {config.DATA.DATASET}. Using Imagenet mean/std." + ) + mean, std = IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD + t.append(transforms.ToTensor()) - t.append(transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)) + t.append(transforms.Normalize(mean, std)) return transforms.Compose(t) diff --git a/data/constants.py b/data/constants.py new file mode 100644 index 00000000..a47b7b89 --- /dev/null +++ b/data/constants.py @@ -0,0 +1,8 @@ +import torch + +data_mean_std = { + "/mnt/10tb/data/inat21/resize-192": ( + torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), + torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), + ), +} diff --git a/data/hierarchical.py b/data/hierarchical.py new file mode 100644 index 00000000..5bcd5057 --- /dev/null +++ b/data/hierarchical.py @@ -0,0 +1,94 @@ +import os + +import torch +from timm.data import Mixup, mixup +from torchvision.datasets import ImageFolder + + +class HierarchicalImageFolder(ImageFolder): + """ + Parses an image folder where the hierarchy is represented as follows: + + 00000_top_middle_..._bottom + 00001_top_middle_..._other + ... + """ + + num_classes = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def find_classes(self, directory): + classes = sorted( + entry.name for entry in os.scandir(directory) if entry.is_dir() + ) + + tier_lookup = {} + class_to_idxs = {} + + for cls in classes: + tiers = make_hierarchical(cls) + + for tier, value in enumerate(tiers): + if tier not in tier_lookup: + tier_lookup[tier] = {} + + if value not in tier_lookup[tier]: + tier_lookup[tier][value] = len(tier_lookup[tier]) + + class_to_idxs[cls] = torch.tensor( + [tier_lookup[tier][value] for tier, value in enumerate(tiers)] + ) + + # Set self.num_classes + self.num_classes = tuple(len(tier) for tier in tier_lookup.values()) + + return classes, class_to_idxs + + +def make_hierarchical(name): + """ + Sometimes the tree is not really a tree; that is, sometimes there are + repeated orders, for example. + + Arguments: + name (str): the complete taxonomic name, separated by '_' + """ + # index is a number + # top is kingdom + index, top, *tiers = name.split("_") + + cleaned = [top] + + complete = top + for tier in tiers: + complete += f"-{tier}" + cleaned.append(complete) + + return cleaned + + +class HierarchicalMixup(Mixup): + def __call__(self, inputs, targets): + assert len(inputs) % 2 == 0, "Batch size should be even when using this" + if self.mode == "elem": + lam = self._mix_elem(inputs) + elif self.mode == "pair": + lam = self._mix_pair(inputs) + else: + lam = self._mix_batch(inputs) + + batch_size, *_ = inputs.shape + assert targets.shape == ( + batch_size, + len(self.num_classes), + ), f"{targets.shape} != {batch_size, len(self.num_classes)}" + + targets = [ + mixup.mixup_target( + target, num_classes, lam, self.label_smoothing, inputs.device + ) + for target, num_classes in zip(targets.T, self.num_classes) + ] + return inputs, targets diff --git a/hierarchical.py b/hierarchical.py new file mode 100644 index 00000000..fdfdf905 --- /dev/null +++ b/hierarchical.py @@ -0,0 +1,77 @@ +import einops +import torch + + +def accuracy(output, target, topk=(1,), hierarchy_level=-1): + """ + Computes the accuracy over the k top predictions for the specified values of k + + Copied from rwightman/pytorch-image-models/timm/utils/metrics.py and modified + to work with hierarchical outputs as well. + + When the output is hierarchical, only returns the accuracy for `hierarchy_level` + (default -1, which is the fine-grained level). + """ + output_levels = 1 + if isinstance(output, list): + output_levels = len(output) + output = output[-1] + + batch_size = output.size(0) + + # Target might have multiple levels because of the hierarchy + if target.squeeze().ndim == 2: + assert target.squeeze().shape == (batch_size, output_levels) + target = target[:, -1] + + maxk = min(max(topk), output.size(1)) + _, pred = output.topk(maxk, dim=1, largest=True, sorted=True) + pred = pred.t() + correct = pred.eq(target.reshape(1, -1).expand_as(pred)) + return [ + correct[: min(k, maxk)].reshape(-1).float().sum(0) * 100.0 / batch_size + for k in topk + ] + + +class FineGrainedCrossEntropyLoss(torch.nn.CrossEntropyLoss): + """ + A cross-entropy used with hierarchical inputs and targets and only + looks at the finest-grained tier (the last level). + """ + + def forward(self, inputs, targets): + fine_grained_inputs = inputs[-1] + fine_grained_targets = targets[:, -1] + return super().forward(fine_grained_inputs, fine_grained_targets) + + +class HierarchicalCrossEntropyLoss(torch.nn.CrossEntropyLoss): + def __init__(self, *args, coeffs=(1.0,), **kwargs): + super().__init__(*args, **kwargs) + + if isinstance(coeffs, torch.Tensor): + coeffs = coeffs.clone().detach().type(torch.float) + else: + coeffs = torch.tensor(coeffs, dtype=torch.float) + + self.register_buffer("coeffs", coeffs) + + def forward(self, inputs, targets): + if not isinstance(targets, list): + targets = einops.rearrange(targets, "batch tiers -> tiers batch") + + assert ( + len(inputs) == len(targets) == len(self.coeffs) + ), f"{len(inputs)} != {len(targets)} != {len(self.coeffs)}" + + losses = torch.stack( + [ + # Need to specify arguments to super() because of some a bug + # with super() in list comprehensions/generators (unclear) + super(HierarchicalCrossEntropyLoss, self).forward(input, target) + for input, target in zip(inputs, targets) + ] + ) + + return torch.dot(self.coeffs, losses) diff --git a/logger.py b/logger.py index d2e0b4b1..642e2167 100644 --- a/logger.py +++ b/logger.py @@ -10,7 +10,9 @@ import os import sys +import torch from termcolor import colored +from torch.utils.tensorboard import SummaryWriter @functools.lru_cache() @@ -46,3 +48,29 @@ def create_logger(output_dir, dist_rank=0, name=""): logger.addHandler(file_handler) return logger + + +class TensorboardWriter: + writer = None + + def __init__(self, output_dir, dist_rank): + if dist_rank == 0: + self.writer = SummaryWriter(log_dir=output_dir) + + def log(self, items, step): + if self.writer is None: + return + + # Copied from huggingface/accelerate/src/accelerate/tracking.py + for k, v in items.items(): + if isinstance(v, (int, float)): + self.writer.add_scalar(k, v, global_step=step) + elif isinstance(v, torch.Tensor): + assert v.numel() == 1 + self.writer.add_scalar(k, v.item(), global_step=step) + elif isinstance(v, str): + self.writer.add_text(k, v, global_step=step) + elif isinstance(v, dict): + self.writer.add_scalars(k, v, global_step=step) + else: + print(f"Can't log {v} because it is {type(v)}!") diff --git a/lr_scheduler.py b/lr_scheduler.py index e991c529..fa0ec4f1 100644 --- a/lr_scheduler.py +++ b/lr_scheduler.py @@ -26,7 +26,7 @@ def build_scheduler(config, optimizer, n_iter_per_epoch): t_initial=(num_steps - warmup_steps) if config.TRAIN.LR_SCHEDULER.WARMUP_PREFIX else num_steps, - t_mul=1.0, + cycle_mul=1.0, lr_min=config.TRAIN.MIN_LR, warmup_lr_init=config.TRAIN.WARMUP_LR, warmup_t=warmup_steps, @@ -47,7 +47,7 @@ def build_scheduler(config, optimizer, n_iter_per_epoch): lr_scheduler = StepLRScheduler( optimizer, decay_t=decay_steps, - decay_rate=config.TRAIN.LR_SCHEDULER.DECAY_RATE, + cycle_decay=config.TRAIN.LR_SCHEDULER.DECAY_RATE, warmup_lr_init=config.TRAIN.WARMUP_LR, warmup_t=warmup_steps, t_in_epochs=False, @@ -61,11 +61,35 @@ def build_scheduler(config, optimizer, n_iter_per_epoch): warmup_t=warmup_steps, t_in_epochs=False, ) + elif config.TRAIN.LR_SCHEDULER.NAME == "constant": + lr_scheduler = ConstantLRScheduler( + optimizer, + warmup_lr_init=config.TRAIN.WARMUP_LR, + warmup_t=warmup_steps, + t_in_epochs=False, + ) return lr_scheduler -class LinearLRScheduler(Scheduler): +class TimmScheduler(Scheduler): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def get_epoch_values(self, epoch: int): + if self.t_in_epochs: + return self._get_lr(epoch) + else: + return None + + def get_update_values(self, num_updates: int): + if not self.t_in_epochs: + return self._get_lr(num_updates) + else: + return None + + +class LinearLRScheduler(TimmScheduler): def __init__( self, optimizer: torch.optim.Optimizer, @@ -115,20 +139,8 @@ def _get_lr(self, t): ] return lrs - def get_epoch_values(self, epoch: int): - if self.t_in_epochs: - return self._get_lr(epoch) - else: - return None - - def get_update_values(self, num_updates: int): - if not self.t_in_epochs: - return self._get_lr(num_updates) - else: - return None - -class MultiStepLRScheduler(Scheduler): +class MultiStepLRScheduler(TimmScheduler): def __init__( self, optimizer: torch.optim.Optimizer, @@ -165,14 +177,31 @@ def _get_lr(self, t): ] return lrs - def get_epoch_values(self, epoch: int): - if self.t_in_epochs: - return self._get_lr(epoch) + +class ConstantLRScheduler(TimmScheduler): + def __init__( + self, + optimizer: torch.optim.Optimizer, + warmup_t=0, + warmup_lr_init=0, + t_in_epochs=True, + ) -> None: + super().__init__(optimizer, param_group_field="lr") + + self.warmup_t = warmup_t + self.warmup_lr_init = warmup_lr_init + self.t_in_epochs = t_in_epochs + if self.warmup_t: + self.warmup_steps = [ + (v - warmup_lr_init) / self.warmup_t for v in self.base_values + ] + super().update_groups(self.warmup_lr_init) else: - return None + self.warmup_steps = [1 for _ in self.base_values] - def get_update_values(self, num_updates: int): - if not self.t_in_epochs: - return self._get_lr(num_updates) + def _get_lr(self, t): + if t < self.warmup_t: + lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps] else: - return None + lrs = [v for v in self.base_values] + return lrs diff --git a/main.py b/main.py index cfe333fc..0793c663 100644 --- a/main.py +++ b/main.py @@ -16,18 +16,24 @@ import torch import torch.backends.cudnn as cudnn import torch.distributed as dist -from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy -from timm.utils import AverageMeter, accuracy +from timm.loss import LabelSmoothingCrossEntropy +from timm.utils import AverageMeter from config import get_config from data import build_loader -from logger import create_logger +from hierarchical import ( + FineGrainedCrossEntropyLoss, + HierarchicalCrossEntropyLoss, + accuracy, +) +from logger import TensorboardWriter, create_logger from lr_scheduler import build_scheduler from models import build_model from optimizer import build_optimizer from utils import ( NativeScalerWithGradNormCount, auto_resume_helper, + batch_size, load_checkpoint, load_pretrained, reduce_tensor, @@ -105,14 +111,6 @@ def parse_option(): "--throughput", action="store_true", help="Test throughput only" ) - # distributed training - parser.add_argument( - "--local_rank", - type=int, - required=True, - help="local rank for DistributedDataParallel", - ) - # for acceleration parser.add_argument( "--fused_window_process", @@ -160,7 +158,12 @@ def main(config): optimizer = build_optimizer(config, model) model = torch.nn.parallel.DistributedDataParallel( - model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False + model, + device_ids=[config.LOCAL_RANK], + broadcast_buffers=False, + find_unused_parameters=False, + gradient_as_bucket_view=True, + static_graph=True, ) loss_scaler = NativeScalerWithGradNormCount() @@ -171,13 +174,22 @@ def main(config): else: lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train)) - if config.AUG.MIXUP > 0.0: - # smoothing is handled with mixup label transform - criterion = SoftTargetCrossEntropy() - elif config.MODEL.LABEL_SMOOTHING > 0.0: + if config.AUG.MIXUP == 0 and config.MODEL.LABEL_SMOOTHING > 0.0: + if config.HIERARHICAL: + raise NotImplementedError( + "We don't support hierarhical loss with label smoothing and no mixup." + ) criterion = LabelSmoothingCrossEntropy(smoothing=config.MODEL.LABEL_SMOOTHING) else: - criterion = torch.nn.CrossEntropyLoss() + # If we have mixup, smoothing is handled with mixup label transform + if config.HIERARHICAL: + criterion = HierarchicalCrossEntropyLoss( + coeffs=config.TRAIN.HIERARCHICAL_COEFFS + ).to(torch.cuda.current_device()) + else: + criterion = torch.nn.CrossEntropyLoss() + + logger.info("Loss function: %s", criterion) max_accuracy = 0.0 @@ -199,7 +211,9 @@ def main(config): max_accuracy = load_checkpoint( config, model_without_ddp, optimizer, lr_scheduler, loss_scaler, logger ) - acc1, acc5, loss = validate(config, data_loader_val, model) + acc1, acc5, loss = validate( + config, data_loader_val, model, config.TRAIN.START_EPOCH - 1 + ) logger.info( f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" ) @@ -208,7 +222,9 @@ def main(config): if config.MODEL.PRETRAINED and (not config.MODEL.RESUME): load_pretrained(config, model_without_ddp, logger) - acc1, acc5, loss = validate(config, data_loader_val, model) + acc1, acc5, loss = validate( + config, data_loader_val, model, config.TRAIN.START_EPOCH - 1 + ) logger.info( f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" ) @@ -247,7 +263,7 @@ def main(config): logger, ) - acc1, acc5, loss = validate(config, data_loader_val, model) + acc1, acc5, loss = validate(config, data_loader_val, model, epoch) logger.info( f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" ) @@ -307,14 +323,20 @@ def train_one_epoch( ) if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0: optimizer.zero_grad() - lr_scheduler.step_update( - (epoch * num_steps + idx) // config.TRAIN.ACCUMULATION_STEPS - ) + if lr_scheduler is not None: + lr_scheduler.step_update( + (epoch * num_steps + idx) // config.TRAIN.ACCUMULATION_STEPS + ) loss_scale_value = loss_scaler.state_dict()["scale"] torch.cuda.synchronize() - loss_meter.update(loss.item(), targets.size(0)) + # We divide by accumulation steps (not sure why) but it makes + # the logged values look weird. So I multiply by it to fix that. + loss_meter.update( + loss.item() * config.TRAIN.ACCUMULATION_STEPS, batch_size(targets) + ) + if grad_norm is not None: # loss_scaler return None if not update norm_meter.update(grad_norm) scaler_meter.update(loss_scale_value) @@ -328,22 +350,44 @@ def train_one_epoch( etas = batch_time.avg * (num_steps - idx) logger.info( f"Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t" - f"eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t wd {wd:.4f}\t" + f"eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t" + f"wd {wd:.4f}\t" f"time {batch_time.val:.4f} ({batch_time.avg:.4f})\t" f"loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t" f"grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f})\t" f"loss_scale {scaler_meter.val:.4f} ({scaler_meter.avg:.4f})\t" f"mem {memory_used:.0f}MB" ) + stats = { + "train_batch_time": batch_time.val, + "train_loss": loss_meter.val, + "train_grad_norm": norm_meter.val, + "train_loss_scale": scaler_meter.val, + "memory_mb": memory_used, + "learning_rate": config.TRAIN.BASE_LR, + } + if lr_scheduler is not None: + stats["learning_rate"] = lr_scheduler.get_update_values( + # Copied from line 326 + (epoch * num_steps + idx) + // config.TRAIN.ACCUMULATION_STEPS + )[0] + + tb_writer.log(stats, (epoch * num_steps + idx)) + epoch_time = time.time() - start logger.info( - f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}" + f"EPOCH {epoch} training took {datetime.timedelta(seconds=int(epoch_time))}" ) + tb_writer.log({"train_time": epoch_time}, epoch) @torch.no_grad() -def validate(config, data_loader, model): - criterion = torch.nn.CrossEntropyLoss() +def validate(config, data_loader, model, epoch): + if config.HIERARHICAL: + criterion = FineGrainedCrossEntropyLoss() + else: + criterion = torch.nn.CrossEntropyLoss() model.eval() batch_time = AverageMeter() @@ -387,6 +431,14 @@ def validate(config, data_loader, model): f"Mem {memory_used:.0f}MB" ) logger.info(f" * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}") + tb_writer.log( + { + "val_acc1": acc1_meter.avg, + "val_acc5": acc5_meter.avg, + "val_loss": loss_meter.avg, + }, + epoch, + ) return acc1_meter.avg, acc5_meter.avg, loss_meter.avg @@ -465,9 +517,10 @@ def throughput(data_loader, model, logger): logger = create_logger( output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" ) + tb_writer = TensorboardWriter(output_dir=config.OUTPUT, dist_rank=dist.get_rank()) if dist.get_rank() == 0: - path = os.path.join(config.OUTPUT, "config.json") + path = os.path.join(config.OUTPUT, "config.yaml") with open(path, "w") as f: f.write(config.dump()) logger.info(f"Full config saved to {path}") diff --git a/models/build.py b/models/build.py index 2d8b07f0..55932d94 100644 --- a/models/build.py +++ b/models/build.py @@ -21,7 +21,7 @@ def build_model(config, is_pretrain=False): import apex as amp layernorm = amp.normalization.FusedLayerNorm - except: + except ImportError: layernorm = None print("To use FusedLayerNorm, please install apex.") else: diff --git a/models/swin_transformer_v2.py b/models/swin_transformer_v2.py index 9d6d2f86..99a7a2cb 100644 --- a/models/swin_transformer_v2.py +++ b/models/swin_transformer_v2.py @@ -13,6 +13,24 @@ from timm.models.layers import DropPath, to_2tuple, trunc_normal_ +class HierarchicalHead(nn.Module): + def __init__(self, num_features, num_classes): + super().__init__() + self.num_classes = tuple(num_classes) + for num_class in self.num_classes: + assert num_class > 0 + + self.heads = nn.ModuleList( + [nn.Linear(num_features, num_class) for num_class in self.num_classes] + ) + + def forward(self, x): + # we do not want to use self.heads(x) because that would feed them through + # each element in the list sequentially, whereas we want x through each head + # individually. + return [head(x) for head in self.heads] + + class Mlp(nn.Module): def __init__( self, @@ -109,6 +127,7 @@ def __init__( self.logit_scale = nn.Parameter( torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True ) + self.register_buffer("logit_clamp_max", torch.log(torch.tensor(1.0 / 0.01))) # mlp to generate continuous relative position bias self.cpb_mlp = nn.Sequential( @@ -200,9 +219,7 @@ def forward(self, x, mask=None): # cosine attention attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) - logit_scale = torch.clamp( - self.logit_scale, max=torch.log(torch.tensor(1.0 / 0.01)) - ).exp() + logit_scale = torch.clamp(self.logit_scale, max=self.logit_clamp_max).exp() attn = attn * logit_scale relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view( @@ -756,11 +773,19 @@ def __init__( self.norm = norm_layer(self.num_features) self.avgpool = nn.AdaptiveAvgPool1d(1) - self.head = ( - nn.Linear(self.num_features, num_classes) - if num_classes > 0 - else nn.Identity() - ) + + # Checks if we are doing hierarchical classification or not. + if isinstance(num_classes, int): + self.head = ( + nn.Linear(self.num_features, num_classes) + if num_classes > 0 + else nn.Identity() + ) + self.hierarchical = False + else: + self.num_classes = tuple(num_classes) + self.head = HierarchicalHead(self.num_features, num_classes) + self.hierarchical = True self.apply(self._init_weights) for bly in self.layers: @@ -805,7 +830,7 @@ def forward(self, x): def flops(self): flops = 0 flops += self.patch_embed.flops() - for i, layer in enumerate(self.layers): + for layer in self.layers: flops += layer.flops() flops += ( self.num_features @@ -813,5 +838,13 @@ def flops(self): * self.patches_resolution[1] // (2**self.num_layers) ) - flops += self.num_features * self.num_classes + if isinstance(self.num_classes, int): + flops += self.num_features * self.num_classes + elif isinstance(self.num_classes, tuple): + for num_class in self.num_classes: + flops += self.num_features * num_class + else: + raise RuntimeError( + f"Internal error: self.num_classes should be int or tuple, not {type(self.num_classes)}" + ) return flops diff --git a/optimizer.py b/optimizer.py index fc1f3167..ddbe6f75 100644 --- a/optimizer.py +++ b/optimizer.py @@ -11,7 +11,7 @@ try: from apex.optimizers import FusedAdam, FusedLAMB -except: +except ImportError: FusedAdam = None FusedLAMB = None print("To use FusedLAMB or FusedAdam, please install apex.") diff --git a/scripts/train.fish b/scripts/train.fish new file mode 100644 index 00000000..b7fcb334 --- /dev/null +++ b/scripts/train.fish @@ -0,0 +1,56 @@ +# Really the only thing we would ever want to change in this script +# are the cfg and the tag. + +function usage + echo "Runs a training script from scratch." + echo + echo "Arguments:" + echo " -h/--help print this message" + echo " --config which config file to use (should be YAML)" + echo " --debug run with only one process rather than 8 so pdb is useable." + echo " --tag tag for the run (I use v0, v1, etc)" + echo " --venv path to the virtual environment (default ./venv/)" +end + +set -l options (fish_opt --short h --long help) +set -a options (fish_opt --short c --long config --required-val --long-only) +set -a options (fish_opt --short t --long tag --required-val --long-only) +set -a options (fish_opt --short v --long venv --long-only) +set -a options (fish_opt --short d --long debug --long-only) + +argparse $options -- $argv + +if not test -z $_flag_help + usage + exit 0 +end + +if test -z $_flag_config + echo "You must provide a --config argument!" + exit 1 +end + +if test -z $_flag_tag + echo "You must provide a --tag argument!" + exit 1 +end + +set venv ./venv +if not test -z $_flag_venv + set venv $_flag_venv +end + +set launcher $venv/bin/torchrun +set launcher_args --nproc_per_node 8 --master_port 12345 + +if not test -z $_flag_debug + set launcher_args --nproc_per_node 1 --master_port 12345 +end + +$launcher $launcher_args \ + main.py \ + --cfg $_flag_config \ + --output runs \ + --tag $_flag_tag \ + --fused_window_process \ + --fused_layernorm diff --git a/utils.py b/utils.py index 8f09bb35..b4b7c79e 100644 --- a/utils.py +++ b/utils.py @@ -203,7 +203,7 @@ def get_grad_norm(parameters, norm_type=2): def auto_resume_helper(output_dir): checkpoints = os.listdir(output_dir) checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith("pth")] - print(f"All checkpoints founded in {output_dir}: {checkpoints}") + print(f"All checkpoints found in {output_dir}: {checkpoints}") if len(checkpoints) > 0: latest_checkpoint = max( [os.path.join(output_dir, d) for d in checkpoints], key=os.path.getmtime @@ -242,6 +242,15 @@ def ampscaler_get_grad_norm(parameters, norm_type: float = 2.0) -> torch.Tensor: return total_norm +def batch_size(tensor_or_list): + if isinstance(tensor_or_list, torch.Tensor): + return tensor_or_list.size(0) + elif isinstance(tensor_or_list, list): + sizes = [tensor.size(0) for tensor in tensor_or_list] + assert all(size == sizes[0] for size in sizes) + return sizes[0] + + class NativeScalerWithGradNormCount: state_dict_key = "amp_scaler" From 105e3116255c40f96bda13746a9f1fc5451d7d04 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Tue, 25 Oct 2022 22:12:04 +0000 Subject: [PATCH 04/47] Add configs --- configs/swinv2/a6000.yaml | 131 ++++++++++++++++++ configs/swinv2/large_inat21_a6000.yaml | 31 +++++ configs/swinv2/large_inat21_constant_lr.yaml | 30 ++++ ...e_patch4_window12_inat21_hierarchical.yaml | 26 ++++ ...window12_inat21_hierarchical_debugging.yml | 40 ++++++ 5 files changed, 258 insertions(+) create mode 100644 configs/swinv2/a6000.yaml create mode 100644 configs/swinv2/large_inat21_a6000.yaml create mode 100644 configs/swinv2/large_inat21_constant_lr.yaml create mode 100644 configs/swinv2/large_patch4_window12_inat21_hierarchical.yaml create mode 100644 configs/swinv2/large_patch4_window12_inat21_hierarchical_debugging.yml diff --git a/configs/swinv2/a6000.yaml b/configs/swinv2/a6000.yaml new file mode 100644 index 00000000..d9930c4a --- /dev/null +++ b/configs/swinv2/a6000.yaml @@ -0,0 +1,131 @@ +AMP_ENABLE: true +AMP_OPT_LEVEL: '' +AUG: + AUTO_AUGMENT: rand-m9-mstd0.5-inc1 + COLOR_JITTER: 0.4 + CUTMIX: 1.0 + CUTMIX_MINMAX: null + MIXUP: 0.8 + MIXUP_MODE: batch + MIXUP_PROB: 1.0 + MIXUP_SWITCH_PROB: 0.5 + RECOUNT: 1 + REMODE: pixel + REPROB: 0.25 +BASE: +- '' +DATA: + BATCH_SIZE: 128 + CACHE_MODE: part + DATASET: inat21 + DATA_PATH: /research/nfs_su_809/cv_datasets/inat21/train_val_224 + IMG_SIZE: 224 + INTERPOLATION: bicubic + NUM_WORKERS: 8 + PIN_MEMORY: true + ZIP_MODE: false +EVAL_MODE: false +LOCAL_RANK: 0 +MODEL: + DROP_PATH_RATE: 0.2 + DROP_RATE: 0.0 + LABEL_SMOOTHING: 0.1 + NAME: swinv2_large_patch4_window7_224_inat21 + NUM_CLASSES: 1000 + PRETRAINED: '' + RESUME: '' + SWIN: + APE: false + DEPTHS: + - 2 + - 2 + - 6 + - 2 + EMBED_DIM: 96 + IN_CHANS: 3 + MLP_RATIO: 4.0 + NUM_HEADS: + - 3 + - 6 + - 12 + - 24 + PATCH_NORM: true + PATCH_SIZE: 4 + QKV_BIAS: true + QK_SCALE: null + WINDOW_SIZE: 7 + SWINV2: + APE: false + DEPTHS: + - 2 + - 2 + - 18 + - 2 + EMBED_DIM: 192 + IN_CHANS: 3 + MLP_RATIO: 4.0 + NUM_HEADS: + - 6 + - 12 + - 24 + - 48 + PATCH_NORM: true + PATCH_SIZE: 4 + PRETRAINED_WINDOW_SIZES: + - 0 + - 0 + - 0 + - 0 + QKV_BIAS: true + WINDOW_SIZE: 7 + SWIN_MLP: + APE: false + DEPTHS: + - 2 + - 2 + - 6 + - 2 + EMBED_DIM: 96 + IN_CHANS: 3 + MLP_RATIO: 4.0 + NUM_HEADS: + - 3 + - 6 + - 12 + - 24 + PATCH_NORM: true + PATCH_SIZE: 4 + WINDOW_SIZE: 7 + TYPE: swinv2 +OUTPUT: /home/stevens.994/projects/Swin-Transformer/model-weights/swinv2_large_patch4_window7_224_inat21/test2 +PRINT_FREQ: 10 +SAVE_FREQ: 1 +SEED: 0 +TAG: test2 +TEST: + CROP: true + SEQUENTIAL: false +THROUGHPUT_MODE: false +TRAIN: + ACCUMULATION_STEPS: 1 + AUTO_RESUME: true + BASE_LR: 0.000125 + CLIP_GRAD: 5.0 + EPOCHS: 180 + LR_SCHEDULER: + DECAY_EPOCHS: 30 + DECAY_RATE: 0.1 + NAME: cosine + MIN_LR: 1.25e-06 + OPTIMIZER: + BETAS: + - 0.9 + - 0.999 + EPS: 1.0e-08 + MOMENTUM: 0.9 + NAME: adamw + START_EPOCH: 0 + USE_CHECKPOINT: false + WARMUP_EPOCHS: 5 + WARMUP_LR: 1.25e-07 + WEIGHT_DECAY: 0.1 diff --git a/configs/swinv2/large_inat21_a6000.yaml b/configs/swinv2/large_inat21_a6000.yaml new file mode 100644 index 00000000..cf448a03 --- /dev/null +++ b/configs/swinv2/large_inat21_a6000.yaml @@ -0,0 +1,31 @@ +DATA: + DATASET: inat21 + IMG_SIZE: 224 + BATCH_SIZE: 32 + DATA_PATH: /mnt/10tb/data/inat21/resize-224 + NUM_WORKERS: 32 +MODEL: + TYPE: swinv2 + NAME: swinv2_large_resize_224_inat21 + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 192 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 6, 12, 24, 48 ] + WINDOW_SIZE: 7 +TRAIN: + EPOCHS: 500 + WARMUP_EPOCHS: 5 + WEIGHT_DECAY: 0.1 + BASE_LR: 1.25e-4 + WARMUP_LR: 1.25e-7 + CLIP_GRAD: 5.0 + # N_GPU * BATCH_SIZE * ACCUMULATION_STEPS = 512 + ACCUMULATION_STEPS: 2 + + LR_SCHEDULER: + NAME: constant + +# Only save checkpoint every 5 epochs +SAVE_FREQ: 5 + diff --git a/configs/swinv2/large_inat21_constant_lr.yaml b/configs/swinv2/large_inat21_constant_lr.yaml new file mode 100644 index 00000000..bfce8b04 --- /dev/null +++ b/configs/swinv2/large_inat21_constant_lr.yaml @@ -0,0 +1,30 @@ +DATA: + DATASET: inat21 + IMG_SIZE: 192 + BATCH_SIZE: 32 + DATA_PATH: /mnt/10tb/data/inat21/resize-192 + NUM_WORKERS: 32 +MODEL: + TYPE: swinv2 + NAME: swinv2_large_resize_192_inat21 + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 192 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 6, 12, 24, 48 ] + WINDOW_SIZE: 12 +TRAIN: + EPOCHS: 240 + WARMUP_EPOCHS: 20 + WEIGHT_DECAY: 0.1 + BASE_LR: 1.0e-3 + WARMUP_LR: 1.0e-7 + CLIP_GRAD: 5.0 + # N_GPU * BATCH_SIZE * ACCUMULATION_STEPS ~= 4096 + ACCUMULATION_STEPS: 16 + + LR_SCHEDULER: + NAME: cosine + +# Only save checkpoint every 5 epochs +SAVE_FREQ: 5 diff --git a/configs/swinv2/large_patch4_window12_inat21_hierarchical.yaml b/configs/swinv2/large_patch4_window12_inat21_hierarchical.yaml new file mode 100644 index 00000000..3d6761bf --- /dev/null +++ b/configs/swinv2/large_patch4_window12_inat21_hierarchical.yaml @@ -0,0 +1,26 @@ +DATA: + DATASET: inat21 + IMG_SIZE: 192 + BATCH_SIZE: 32 + DATA_PATH: /mnt/10tb/data/inat21/resize-192 + NUM_WORKERS: 32 +MODEL: + TYPE: swinv2 + NAME: swinv2_large_patch4_window12_resize_192_inat21 + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 192 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 6, 12, 24, 48 ] + WINDOW_SIZE: 12 +TRAIN: + EPOCHS: 90 + WARMUP_EPOCHS: 5 + WEIGHT_DECAY: 0.1 + BASE_LR: 1.25e-4 # 4096 batch-size + WARMUP_LR: 1.25e-7 + MIN_LR: 1.25e-6 + HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] + # We want N_GPU * BATCH_SIZE * ACCUMULATION_STEPS ~= 1024 + ACCUMULATION_STEPS: 4 +HIERARHICAL: true diff --git a/configs/swinv2/large_patch4_window12_inat21_hierarchical_debugging.yml b/configs/swinv2/large_patch4_window12_inat21_hierarchical_debugging.yml new file mode 100644 index 00000000..9462ab9d --- /dev/null +++ b/configs/swinv2/large_patch4_window12_inat21_hierarchical_debugging.yml @@ -0,0 +1,40 @@ +DATA: + DATASET: inat21 + IMG_SIZE: 192 + BATCH_SIZE: 32 + DATA_PATH: /mnt/10tb/data/inat21/resize-192 + NUM_WORKERS: 32 +AUG: + COLOR_JITTER: 0.0 + AUTO_AUGMENT: "none" + REPROB: 0.0 + RECOUNT: 0 + MIXUP: 0.0 + CUTMIX: 0.0 + MIXUP_PROB: 0.0 +MODEL: + TYPE: swinv2 + NAME: swinv2_large_patch4_window12_resize_192_inat21_debugging + SWINV2: + EMBED_DIM: 192 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 6, 12, 24, 48 ] + WINDOW_SIZE: 12 + # Regularization + DROP_RATE: 0.0 + DROP_PATH_RATE: 0.0 + LABEL_SMOOTHING: 0.0 +TRAIN: + OVERFIT_BATCHES: 8 + EPOCHS: 90 + WARMUP_EPOCHS: 0 + WEIGHT_DECAY: 0.0 + BASE_LR: 1.25e-4 # 4096 batch-size + WARMUP_LR: 1.25e-4 + MIN_LR: 1.25e-4 + HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] + # We want N_GPU * BATCH_SIZE * ACCUMULATION_STEPS ~= 1024 + ACCUMULATION_STEPS: 4 + LR_SCHEDULER: + NAME: "none" +HIERARHICAL: true From 018e6ec58b2a95c086b3d24652f5c63c552e033e Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Tue, 25 Oct 2022 22:12:10 +0000 Subject: [PATCH 05/47] Add inat dataset helpers --- data/inat/__init__.py | 0 data/inat/__main__.py | 70 +++++++++++++++ data/inat/datasets.py | 196 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 266 insertions(+) create mode 100644 data/inat/__init__.py create mode 100644 data/inat/__main__.py create mode 100644 data/inat/datasets.py diff --git a/data/inat/__init__.py b/data/inat/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/data/inat/__main__.py b/data/inat/__main__.py new file mode 100644 index 00000000..b9824310 --- /dev/null +++ b/data/inat/__main__.py @@ -0,0 +1,70 @@ +import argparse + +from .. import hierarchical +from . import datasets + + +def preprocess_cli(args): + datasets.preprocess_dataset(args.root, args.stage, args.strategy, args.size) + + +def normalize_cli(args): + std, mean = datasets.load_statistics(args.directory, use_cache=False) + print("Add this to a constants.py file:") + print( + f""" +"{args.directory}": ( + torch.tensor({std.tolist()}), + torch.tensor({mean.tolist()}), +),""" + ) + + +def num_classes_cli(args): + dataset = hierarchical.HierarchicalImageFolder( + args.directory, + ) + + print("Add this to your config .yml file to the model section:") + print(f"num_classes: {list(dataset.num_classes)}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(help="Available commands.") + + # Preprocess + preprocess_parser = subparsers.add_parser("preprocess", help="Preprocess data.") + preprocess_parser.add_argument( + "root", + help="Root data folder. Should contain folders compressed/ and raw/train, raw/val, etc.", + ) + preprocess_parser.add_argument( + "stage", choices=datasets.Inat21.stages, help="Data stage to preprocess" + ) + preprocess_parser.add_argument( + "strategy", choices=datasets.Inat21.strategies, help="Preprocessing strategy" + ) + preprocess_parser.add_argument("size", type=int, help="Image size in pixels") + preprocess_parser.set_defaults(func=preprocess_cli) + + # Normalize + normalize_parser = subparsers.add_parser( + "normalize", help="Measure std. dev. and mean of dataset." + ) + normalize_parser.add_argument("directory", help="Data folder") + normalize_parser.set_defaults(func=normalize_cli) + + # Number of classes + num_classes_parser = subparsers.add_parser( + "num-classes", help="Measure number of classes in dataset." + ) + num_classes_parser.add_argument("directory", help="Data folder") + num_classes_parser.set_defaults(func=num_classes_cli) + + args = parser.parse_args() + + if hasattr(args, "func"): + args.func(args) + else: + parser.print_usage() diff --git a/data/inat/datasets.py b/data/inat/datasets.py new file mode 100644 index 00000000..aaab5691 --- /dev/null +++ b/data/inat/datasets.py @@ -0,0 +1,196 @@ +import concurrent.futures +import os +import pathlib +import warnings + +import cv2 +import einops +import timm.data +import torch +import torchvision +from tqdm.auto import tqdm + + +def load_statistics(directory): + """ + Need to calculate mean and std for the individual channels so we can normalize the images. + """ + + # 1. Load the images into a dataset, then a dataloader so we can iterate over them quickly. + dataset = timm.data.ImageDataset( + root=directory, transform=torchvision.transforms.ToTensor() + ) + channels, height, width = dataset[0][0].shape + + total = torch.zeros((channels,)) + total_squared = torch.zeros((channels,)) + + dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, num_workers=32) + + for batch, _ in tqdm(dataloader): + total += einops.reduce(batch, "batch channel height width -> channel", "sum") + total_squared += einops.reduce( + torch.mul(batch, batch), "batch channel height width -> channel", "sum" + ) + + divisor = len(dataset) * width * height + + mean = total / divisor + var = total_squared / divisor - torch.mul(mean, mean) + std = torch.sqrt(var) + + return std, mean + + +class Inat21: + val_tar_gz_hash = "f6f6e0e242e3d4c9569ba56400938afc" + train_tar_gz_hash = "e0526d53c7f7b2e3167b2b43bb2690ed" + train_mini_tar_gz_hash = "db6ed8330e634445efc8fec83ae81442" + strategies = ("resize", "pad") + stages = ("train", "val", "train_mini") + + def __init__(self, root: str, stage: str, strategy: str, size: int): + self.root = root + self.stage = stage + self._check_stage(stage) + + self.strategy = strategy + self._check_strategy(strategy) + + self.size = size + self._check_size(size) + + @property + def directory(self) -> str: + return os.path.join(self.root, f"{self.strategy}-{self.size}", self.stage) + + @property + def ffcv(self) -> str: + return os.path.join( + self.root, f"{self.stage}-{self.strategy}-{self.size}.beton" + ) + + @property + def tar_file(self): + return os.path.join(self.root, "compressed", f"{self.stage}.tar.gz") + + @property + def raw_dir(self): + return os.path.join(self.root, "raw", self.stage) + + def _check_stage(self, stage): + if stage not in self.stages: + raise ValueError(f"Stage '{stage}' must be one of {self.stages}") + + def _check_strategy(self, strategy): + if strategy not in self.strategies: + raise ValueError(f"Strategy '{strategy}' must be one of {self.strategies}") + + def _check_size(self, size): + if not isinstance(size, int): + raise ValueError(f"Size {size} must be int; not {type(int)}") + + def check(self): + # If /.finished doesn't exist, we need to preprocess. + if not os.path.isfile(finished_file_path(self.directory)): + warnings.warn( + f"Data not processed in {self.directory}! " + "You should run:\n\n" + f"\tpython -m src.data preprocess {self.root} {self.stage} {self.strategy} {self.size}" + "\n\nAnd then run this script again!" + ) + raise RuntimeError(f"Data {self.directory} not pre-processed!") + + +def finished_file_path(directory): + return os.path.join(directory, ".finished") + + +def preprocess_class(cls_dir, output_dir, strategy, size): + cls = os.path.basename(cls_dir) + output_dir = os.path.join(output_dir, cls) + os.makedirs(output_dir, exist_ok=True) + + with os.scandir(cls_dir) as entries: + for entry in entries: + if not entry.is_file(): + continue + + im = cv2.imread(entry.path) + im = preprocess_image(im, strategy, size) + output_path = os.path.join(output_dir, entry.name) + if not cv2.imwrite(output_path, im): + raise RuntimeError(output_path) + + +def preprocess_image(im, strategy, size): + if strategy == "resize": + return cv2.resize(im, (size, size), interpolation=cv2.INTER_LINEAR) + elif strategy == "pad": + # https://stackoverflow.com/questions/43391205/add-padding-to-images-to-get-them-into-the-same-shape + raise NotImplementedError() + else: + raise NotImplementedError() + + +def parent_of(path: str): + return pathlib.Path(path).parents[0] + + +def preprocess_dataset(root: str, stage: str, strategy: str, size: int) -> None: + inat = Inat21(root, stage, strategy, size) + + err_msg = ( + f"Can't prepare data for stage {stage}, strategy {strategy} and size {size}." + ) + + # 1. If the directory does not exist, ask the user to fix that for us. + if not os.path.isdir(inat.raw_dir): + # Check that the tar exists + if not os.path.isfile(inat.tar_file): + warn_msg = f"Please download the appropriate .tar.gz file to {root} for stage {stage}." + if "raw" in root: + warn_msg += f"\n\nYour root path should contain a 'raw' directory; did you mean to use {parent_of(root)}?\n" + elif "compressed" in root: + warn_msg += f"\n\nYour root path should contain a 'compressed' directory; did you mean to use {parent_of(root)}?\n" + + warnings.warn(warn_msg) + + raise RuntimeError(err_msg) + else: + warnings.warn( + f"Please untar {inat.tar_file} in {root}. Probably need to run 'cd {root}; tar -xvf {inat.tar_file}" + ) + raise RuntimeError(err_msg) + + # 2. Now that we know the raw directory exists, we need to convert it + # to a processed directory + out_path = os.path.join(root, inat.directory) + + # 3. Make sure the directory exists + if not os.path.isdir(out_path): + os.makedirs(out_path) + + # 4. For all raw files, process and save to directory. We do this with + # a process pool because it is both I/O (read/write) and CPU (image processing) + # bound. + with os.scandir(inat.raw_dir) as entries: + directories = [entry.path for entry in entries if entry.is_dir()] + # print(directories) + print(f"Found {len(directories)} directories to preprocess.") + + with concurrent.futures.ThreadPoolExecutor() as executor: + futures = [ + executor.submit(preprocess_class, directory, out_path, strategy, size) + for directory in tqdm(directories) + ] + + print(f"Submitted {len(futures)} jobs to executor.") + + for future in tqdm( + concurrent.futures.as_completed(futures), total=len(futures) + ): + future.result() + + # 5. Save a sentinel file called .finished + open(finished_file_path(out_path), "w").close() From e21edeb17aaf693810316d1f39049ed66876014c Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Tue, 25 Oct 2022 22:12:20 +0000 Subject: [PATCH 06/47] Add script to parse swin log files --- scripts/parse_logs.py | 302 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 302 insertions(+) create mode 100644 scripts/parse_logs.py diff --git a/scripts/parse_logs.py b/scripts/parse_logs.py new file mode 100644 index 00000000..5815c863 --- /dev/null +++ b/scripts/parse_logs.py @@ -0,0 +1,302 @@ +""" +This scripts parses the training logs to graph both training loss and validation accuracy over time. +""" + +import argparse +import dataclasses +import re + +import matplotlib.pyplot as plt +import preface + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("file", help="Log file to parse. Typically named log_rank0.txt") + parser.add_argument( + "--last", help="How many of the latest epochs to look at.", default=10, type=int + ) + + return parser.parse_args() + + +@dataclasses.dataclass +class ValidationLine: + epoch: int + batch: int + batch_max: int + loss: float + mean_loss: float + acc1: float + mean_acc1: float + acc5: float + mean_acc5: float + + @classmethod + def from_raw_line(cls, line, last_train): + if "Test" not in line: + return None + + # Example line: + # [2022-06-08 07:34:50 swinv2_large_patch4_window7_224_inat21](main.py 258): INFO Test: [0/196] Time 1.895 (1.895) Loss 0.8066 (0.8066) Acc@1 84.766 (84.766) Acc@5 95.312 (95.312) Mem 36916MB + + pattern = r""" +^\[.*?\]\ +\(main.py\ \d+\):\ +INFO\ Test:\ +\[(?P\d+)/(?P\d+)\] +\t +Time\ \d+.\d+\ \(\d+.\d+\) +\t +Loss\ (?P[\w.]+)\ \((?P[\w.]+)\) +\t +Acc@1\ (?P[\w.]+)\ \((?P[\w.]+)\) +\t +Acc@5\ (?P[\w.]+)\ \((?P[\w.]+)\) +\t +Mem\ (?P.*) +$ +""" + + match = re.match(pattern, line, re.VERBOSE) + + if not match: + print(repr(line)) + return None + + epoch = 0 + if last_train: + epoch = last_train.epoch + + return cls( + epoch, + int(match.group("batch")), + int(match.group("batch_max")), + float(match.group("loss")), + float(match.group("mean_loss")), + float(match.group("acc1")), + float(match.group("mean_acc1")), + float(match.group("acc5")), + float(match.group("mean_acc5")), + ) + + +@dataclasses.dataclass +class TrainLine: + epoch: int + epoch_max: int + batch: int + batch_max: int + lr: float + wd: float + loss: float + mean_loss: float + grad_norm: float + mean_grad_norm: float + loss_scale: float + mean_loss_scale: float + + @classmethod + def from_raw_line(cls, line): + if "Train" not in line: + return None + + # Example line: + # [2022-06-07 20:54:56 swinv2_large_patch4_window7_224_inat21] (main.py 209): INFO Train: [56/90][700/5247]\teta 3:11:57 lr 0.000040\t wd 0.1000\ttime 2.5379 (2.5330)\tloss 4.3632 (3.7333)\tgrad_norm 9.5996 (inf)\tloss_scale 8192.0000 (8542.5849)\tmem 36916MB + + pattern = r""" +^\[.*?\]\ # [2022-06-08 08:35:04 ...] +\(main.py\ \d+\):\ # +INFO\ Train:\ # +\[(?P\d+)/(?P\d+)\] # [56/90] +\[(?P\d+)/(?P\d+)\] # [700/5247] +\t # +eta\ (\d\ day,\ )?\d+:\d\d:\d\d # eta 3:11:57 +\ # +lr\ (?P\d\.\d+) # lr 0.000040 +\t # +wd\ (?P\d\.\d+) # wd 0.1000 +\t # +time\ \d+\.\d+\ \(\d+\.\d+\) # time 2.5379 +\t # +loss\ (?P[\w.]+)\ \((?P[\w.]+)\) # loss 4.3632 (3.7333) +\t # +grad_norm\ (?P[\w.]+)\ \((?P[\w.]+)\) # grad_norm 9.5996 (inf) +\t # +loss_scale\ (?P[\w.]+)\ \((?P[\w.]+)\) # loss_scale 8192.0000 (8542.5849) +\t # +mem\ (?P.*) # mem 36916MB +$ +""" + + match = re.match(pattern, line, re.VERBOSE) + + if not match: + print(repr(line)) + return None + + return cls( + int(match.group("epoch")), + int(match.group("epoch_max")), + int(match.group("batch")), + int(match.group("batch_max")), + float(match.group("lr")), + float(match.group("wd")), + float(match.group("loss")), + float(match.group("mean_loss")), + float(match.group("grad_norm")), + float(match.group("mean_grad_norm")), + float(match.group("loss_scale")), + float(match.group("mean_loss_scale")), + ) + + +def parse_file(filepath): + """ + Turns each line into a TrainLine instance. + """ + with open(filepath) as fd: + raw = [line.strip() for line in fd] + + lines = [] + last_train_line = None + for line in raw: + train_line = TrainLine.from_raw_line(line) + if train_line: + last_train_line = train_line + + val_line = ValidationLine.from_raw_line(line, last_train_line) + + lines.append(train_line or val_line) + + # Filter Nones + return [line for line in lines if line] + + +def filter_same_epochs(lines): + """ + If any of the lines have the same epoch (from training restarting), use only the most recent one. + """ + + if len(set(line.epoch for line in lines)) == len(lines): + return lines + + assert sorted(lines, key=lambda line: line.epoch) == lines + + cleaned = [] + for second, first in preface.grouped(list(reversed(lines)), size=2): + if first.epoch == second.epoch: + continue + + cleaned.append(second) + + assert len(set(line.epoch for line in cleaned)) == len(cleaned) + + return list(reversed(cleaned)) + + +def plot_losses(train_epochs, train_loss, val_epochs, val_loss): + fig, ax = plt.subplots() + + ax.plot(train_epochs, train_loss, color="tab:blue", label="Train Loss") + ax.plot(val_epochs, val_loss, color="tab:orange", label="Val. Loss") + ax.set_xlabel("Epochs") + ax.set_ylabel("Mean ") + + fig.legend() + + return fig + + +def visualize(data, run_name, last): + # Visualize validation error per epoch + val_lines = [line for line in data if isinstance(line, ValidationLine)] + + true_batch_max = max(line.batch for line in val_lines) + last_batches = [line for line in val_lines if line.batch == true_batch_max] + last_batches = filter_same_epochs(last_batches) + + val_loss = [line.mean_loss for line in last_batches] + val_epochs = [line.epoch for line in last_batches] + assert len(set(val_epochs)) == len(val_epochs) + + # Visualize training loss per epoch. + # Use mean_loss for the last batch in each epoch. + train_lines = [line for line in data if isinstance(line, TrainLine)] + + true_batch_max = max(line.batch for line in train_lines) + last_batches = [ + line + for line in train_lines + if line.batch == true_batch_max and line.epoch in val_epochs + ] + last_batches = filter_same_epochs(last_batches) + + train_loss = [line.mean_loss for line in last_batches] + train_epochs = [line.epoch for line in last_batches] + assert len(set(train_epochs)) == len(train_epochs) + + assert val_epochs == train_epochs + + fig = plot_losses(train_epochs, train_loss, val_epochs, val_loss) + fig.suptitle(f"{run_name} Progress") + fig.savefig("training-loss.pdf") + + # Also look at the latest values + + def clip(lst): + return lst[-last:] + + fig = plot_losses( + clip(train_epochs), clip(train_loss), clip(val_epochs), clip(val_loss) + ) + + fig.suptitle(f"{run_name} Last {last} Epochs") + fig.savefig(f"training-loss-{last}-epochs.pdf") + + +def get_run_name(file): + pattern = r".*?/.*?/(.*?)/.*\.txt" + + return re.match(pattern, file).group(1) + + +def plot_lr(data): + import matplotlib.pyplot as plt + + train_lines = [line for line in data if isinstance(line, TrainLine)] + ys = [line.lr for line in train_lines] + xs = list(range(len(ys))) + + fig, ax = plt.subplots() + + ax.plot(xs, ys, linewidth=0.1) + + fig.savefig("learning-rates.pdf") + + +def best_validation_epoch(data): + # Visualize validation error per epoch + val_lines = [line for line in data if isinstance(line, ValidationLine)] + + true_batch_max = max(line.batch for line in val_lines) + last_batches = [line for line in val_lines if line.batch == true_batch_max] + last_batches = filter_same_epochs(last_batches) + + return min(last_batches, key=lambda l: l.mean_loss) + + +def main(): + args = parse_args() + + data = parse_file(args.file) + + visualize(data, get_run_name(args.file), last=args.last) + + plot_lr(data) + + print(best_validation_epoch(data)) + + +if __name__ == "__main__": + main() From 717247935767091e593d0fd8dea9c1485c422813 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Tue, 25 Oct 2022 22:12:42 +0000 Subject: [PATCH 07/47] Add hparams to tensorboard logging --- logger.py | 10 +++++++++- main.py | 2 ++ utils.py | 24 ++++++++++++++++++++++++ 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/logger.py b/logger.py index 642e2167..727bc2fc 100644 --- a/logger.py +++ b/logger.py @@ -54,8 +54,16 @@ class TensorboardWriter: writer = None def __init__(self, output_dir, dist_rank): + self.output_dir = output_dir + if dist_rank == 0: - self.writer = SummaryWriter(log_dir=output_dir) + self.writer = SummaryWriter(log_dir=self.output_dir) + + def add_hparams(self, hparams, metrics): + if self.writer is None: + return + + self.writer.add_hparams(hparams, metrics, run_name=self.output_dir) def log(self, items, step): if self.writer is None: diff --git a/main.py b/main.py index 0793c663..390e02c3 100644 --- a/main.py +++ b/main.py @@ -38,6 +38,7 @@ load_pretrained, reduce_tensor, save_checkpoint, + to_hparams, ) @@ -380,6 +381,7 @@ def train_one_epoch( f"EPOCH {epoch} training took {datetime.timedelta(seconds=int(epoch_time))}" ) tb_writer.log({"train_time": epoch_time}, epoch) + tb_writer.add_hparams(to_hparams(config), stats) @torch.no_grad() diff --git a/utils.py b/utils.py index b4b7c79e..29c664d8 100644 --- a/utils.py +++ b/utils.py @@ -7,6 +7,7 @@ import os +import preface.dict import torch import torch.distributed as dist from torch._six import inf @@ -288,3 +289,26 @@ def state_dict(self): def load_state_dict(self, state_dict): self._scaler.load_state_dict(state_dict) + + +def to_hparams(cfg): + hparams = {} + for key, value in cfg.items(): + key = key.lower() + if ( + isinstance(value, int) + or isinstance(value, float) + or isinstance(value, str) + or isinstance(value, bool) + or isinstance(value, torch.Tensor) + or value is None + ): + hparams[key] = value + elif isinstance(value, dict): + hparams[key] = to_hparams(value) + elif isinstance(value, list) or isinstance(value, tuple): + hparams[key] = "[" + ", ".join(repr(v) for v in value) + "]" + else: + raise ValueError(f"Don't know how to handle {key}: {value}!") + + return preface.dict.flattened(hparams) From 8f6c3a040d025b1d3bc78535d2562b8ac16f3a1a Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Tue, 25 Oct 2022 23:11:03 +0000 Subject: [PATCH 08/47] add pyenv version file --- .python-version | 1 + 1 file changed, 1 insertion(+) create mode 100644 .python-version diff --git a/.python-version b/.python-version new file mode 100644 index 00000000..1281604a --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.10.7 From 61589b8126753196e2737e4babd5dadbd44305fb Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 27 Oct 2022 17:07:35 +0000 Subject: [PATCH 09/47] Update data normalization practices --- data/build.py | 5 ++--- data/constants.py | 4 ++++ data/inat/__main__.py | 2 +- data/inat/datasets.py | 4 +--- 4 files changed, 8 insertions(+), 7 deletions(-) diff --git a/data/build.py b/data/build.py index 6b6d3b0c..7fd74f27 100644 --- a/data/build.py +++ b/data/build.py @@ -220,10 +220,9 @@ def build_transform(is_train, config): elif config.DATA.DATASET in data_mean_std: mean, std = data_mean_std[config.DATA.DATASET] else: - print( - "Can't find mean/std for {config.DATA.DATASET} at {config.DATA.DATASET}. Using Imagenet mean/std." + raise RuntimeError( + f"Can't find mean/std for {config.DATA.DATASET} at {config.DATA.DATASET}. Please add it to data/constants.py (try using python -m data.inat normalize for iNat)." ) - mean, std = IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD t.append(transforms.ToTensor()) t.append(transforms.Normalize(mean, std)) diff --git a/data/constants.py b/data/constants.py index a47b7b89..3919bdc4 100644 --- a/data/constants.py +++ b/data/constants.py @@ -5,4 +5,8 @@ torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), ), + "/mnt/10tb/data/inat21/resize-224": ( + torch.tensor([0.23762744665145874, 0.2292044311761856, 0.24757201969623566]), + torch.tensor([0.4632636606693268, 0.48004215955734253, 0.37622377276420593]), + ), } diff --git a/data/inat/__main__.py b/data/inat/__main__.py index b9824310..0feeb2df 100644 --- a/data/inat/__main__.py +++ b/data/inat/__main__.py @@ -9,7 +9,7 @@ def preprocess_cli(args): def normalize_cli(args): - std, mean = datasets.load_statistics(args.directory, use_cache=False) + std, mean = datasets.load_statistics(args.directory) print("Add this to a constants.py file:") print( f""" diff --git a/data/inat/datasets.py b/data/inat/datasets.py index aaab5691..dcf2e356 100644 --- a/data/inat/datasets.py +++ b/data/inat/datasets.py @@ -15,8 +15,6 @@ def load_statistics(directory): """ Need to calculate mean and std for the individual channels so we can normalize the images. """ - - # 1. Load the images into a dataset, then a dataloader so we can iterate over them quickly. dataset = timm.data.ImageDataset( root=directory, transform=torchvision.transforms.ToTensor() ) @@ -25,7 +23,7 @@ def load_statistics(directory): total = torch.zeros((channels,)) total_squared = torch.zeros((channels,)) - dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, num_workers=32) + dataloader = torch.utils.data.DataLoader(dataset, batch_size=256, num_workers=32) for batch, _ in tqdm(dataloader): total += einops.reduce(batch, "batch channel height width -> channel", "sum") From 1d0f3edf4bafbfd1336c1edbb9ee7eef76754099 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 27 Oct 2022 17:07:40 +0000 Subject: [PATCH 10/47] Update hparams logging --- logger.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/logger.py b/logger.py index 727bc2fc..49d44722 100644 --- a/logger.py +++ b/logger.py @@ -13,6 +13,7 @@ import torch from termcolor import colored from torch.utils.tensorboard import SummaryWriter +from torch.utils.tensorboard.summary import hparams @functools.lru_cache() @@ -57,13 +58,13 @@ def __init__(self, output_dir, dist_rank): self.output_dir = output_dir if dist_rank == 0: - self.writer = SummaryWriter(log_dir=self.output_dir) + self.writer = CorrectedSummaryWriter(log_dir=self.output_dir) def add_hparams(self, hparams, metrics): if self.writer is None: return - self.writer.add_hparams(hparams, metrics, run_name=self.output_dir) + self.writer.add_hparams(hparams, metrics, run_name="") def log(self, items, step): if self.writer is None: @@ -82,3 +83,17 @@ def log(self, items, step): self.writer.add_scalars(k, v, global_step=step) else: print(f"Can't log {v} because it is {type(v)}!") + + +class CorrectedSummaryWriter(SummaryWriter): + def add_hparams(self, hparam_dict, metric_dict): + torch._C._log_api_usage_once("tensorboard.logging.add_hparams") + if type(hparam_dict) is not dict or type(metric_dict) is not dict: + raise TypeError("hparam_dict and metric_dict should be dictionary.") + exp, ssi, sei = hparams(hparam_dict, metric_dict) + + self.file_writer.add_summary(exp) + self.file_writer.add_summary(ssi) + self.file_writer.add_summary(sei) + for k, v in metric_dict.items(): + self.add_scalar(k, v) From 35f1ce27bb1079cc2b8963c5b54a0f172de962bc Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Wed, 26 Oct 2022 20:33:44 +0000 Subject: [PATCH 11/47] Spelling error --- config.py | 2 +- data/build.py | 4 ++-- main.py | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/config.py b/config.py index aafdfdf0..6a86ccbc 100644 --- a/config.py +++ b/config.py @@ -243,7 +243,7 @@ # ----------------------------------------------------------------------------- # Whether we are doing hierarchical classification -_C.HIERARHICAL = False +_C.HIERARCHICAL = False # [SimMIM] Whether to enable pytorch amp, overwritten by command line argument _C.ENABLE_AMP = False diff --git a/data/build.py b/data/build.py index 7fd74f27..e261f454 100644 --- a/data/build.py +++ b/data/build.py @@ -117,7 +117,7 @@ def build_loader(config): label_smoothing=config.MODEL.LABEL_SMOOTHING, num_classes=config.MODEL.NUM_CLASSES, ) - if config.HIERARHICAL: + if config.HIERARCHICAL: mixup_fn = HierarchicalMixup(**mixup_args) else: mixup_fn = Mixup(**mixup_args) @@ -158,7 +158,7 @@ def build_dataset(is_train, config): prefix = "train" if is_train else "val" root = os.path.join(config.DATA.DATA_PATH, prefix) - if config.HIERARHICAL: + if config.HIERARCHICAL: dataset = HierarchicalImageFolder(root, transform=transform) nb_classes = dataset.num_classes else: diff --git a/main.py b/main.py index 390e02c3..98ac6a3e 100644 --- a/main.py +++ b/main.py @@ -176,14 +176,14 @@ def main(config): lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train)) if config.AUG.MIXUP == 0 and config.MODEL.LABEL_SMOOTHING > 0.0: - if config.HIERARHICAL: + if config.HIERARCHICAL: raise NotImplementedError( "We don't support hierarhical loss with label smoothing and no mixup." ) criterion = LabelSmoothingCrossEntropy(smoothing=config.MODEL.LABEL_SMOOTHING) else: # If we have mixup, smoothing is handled with mixup label transform - if config.HIERARHICAL: + if config.HIERARCHICAL: criterion = HierarchicalCrossEntropyLoss( coeffs=config.TRAIN.HIERARCHICAL_COEFFS ).to(torch.cuda.current_device()) @@ -386,7 +386,7 @@ def train_one_epoch( @torch.no_grad() def validate(config, data_loader, model, epoch): - if config.HIERARHICAL: + if config.HIERARCHICAL: criterion = FineGrainedCrossEntropyLoss() else: criterion = torch.nn.CrossEntropyLoss() From 5aeec4f11d7816c97697bdeb3b8c124f5976abc3 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Thu, 27 Oct 2022 20:58:02 +0000 Subject: [PATCH 12/47] README instructions --- README.md | 339 ++++++++---------------------------------------------- 1 file changed, 51 insertions(+), 288 deletions(-) diff --git a/README.md b/README.md index 816bf731..a2dc800c 100644 --- a/README.md +++ b/README.md @@ -1,310 +1,73 @@ # Swin Transformer -[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/swin-transformer-v2-scaling-up-capacity-and/object-detection-on-coco)](https://paperswithcode.com/sota/object-detection-on-coco?p=swin-transformer-v2-scaling-up-capacity-and) -[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/swin-transformer-v2-scaling-up-capacity-and/instance-segmentation-on-coco)](https://paperswithcode.com/sota/instance-segmentation-on-coco?p=swin-transformer-v2-scaling-up-capacity-and) -[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/swin-transformer-v2-scaling-up-capacity-and/semantic-segmentation-on-ade20k)](https://paperswithcode.com/sota/semantic-segmentation-on-ade20k?p=swin-transformer-v2-scaling-up-capacity-and) -[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/swin-transformer-v2-scaling-up-capacity-and/action-classification-on-kinetics-400)](https://paperswithcode.com/sota/action-classification-on-kinetics-400?p=swin-transformer-v2-scaling-up-capacity-and) +[Link to original Swin Transformer project](https://github.com/microsoft/Swin-Transformer) -This repo is the official implementation of ["Swin Transformer: Hierarchical Vision Transformer using Shifted Windows"](https://arxiv.org/pdf/2103.14030.pdf) as well as the follow-ups. It currently includes code and models for the following tasks: +## Installation Instructions -> **Image Classification**: Included in this repo. See [get_started.md](get_started.md) for a quick start. +1. Set up python packages -> **Object Detection and Instance Segmentation**: See [Swin Transformer for Object Detection](https://github.com/SwinTransformer/Swin-Transformer-Object-Detection). - -> **Semantic Segmentation**: See [Swin Transformer for Semantic Segmentation](https://github.com/SwinTransformer/Swin-Transformer-Semantic-Segmentation). - -> **Video Action Recognition**: See [Video Swin Transformer](https://github.com/SwinTransformer/Video-Swin-Transformer). - -> **Semi-Supervised Object Detection**: See [Soft Teacher](https://github.com/microsoft/SoftTeacher). - -> **SSL: Contrasitive Learning**: See [Transformer-SSL](https://github.com/SwinTransformer/Transformer-SSL). - -> **SSL: Masked Image Modeling**: See [get_started.md#simmim-support](https://github.com/microsoft/Swin-Transformer/blob/main/get_started.md#simmim-support). - -> **Mixture-of-Experts**: See [get_started](get_started.md#mixture-of-experts-support) for more instructions. - -> **Feature-Distillation**: Will appear in [Feature-Distillation](https://github.com/SwinTransformer/Feature-Distillation). - -## Activity notification - -* 09/18/2022: Organizing ECCV Workshop [*Computer Vision in the Wild (CVinW)*](https://computer-vision-in-the-wild.github.io/eccv-2022/), where two challenges are hosted to evaluate the zero-shot, few-shot and full-shot performance of pre-trained vision models in downstream tasks: - - [``*Image Classification in the Wild (ICinW)*''](https://eval.ai/web/challenges/challenge-page/1832/overview) Challenge evaluates on 20 image classification tasks. - - [``*Object Detection in the Wild (ODinW)*''](https://eval.ai/web/challenges/challenge-page/1839/overview) Challenge evaluates on 35 object detection tasks. - - -$\qquad$ [ [Workshop]](https://computer-vision-in-the-wild.github.io/eccv-2022/) $\qquad$ [ [IC Challenge] ](https://eval.ai/web/challenges/challenge-page/1832/overview) -$\qquad$ [ [OD Challenge] ](https://eval.ai/web/challenges/challenge-page/1839/overview) - -## Updates - -***09/24/2022*** - -1. Merged [SimMIM](https://github.com/microsoft/SimMIM), which is a **Masked Image Modeling** based pre-training approach applicable to Swin and SwinV2 (and also applicable for ViT and ResNet). Please refer to [get started with SimMIM](get_started.md#simmim-support) to play with SimMIM pre-training. - -2. Released a series of Swin and SwinV2 models pre-trained using the SimMIM approach (see [MODELHUB for SimMIM](MODELHUB.md#simmim-pretrained-swin-v2-models)), with model size ranging from SwinV2-Small-50M to SwinV2-giant-1B, data size ranging from ImageNet-1K-10% to ImageNet-22K, and iterations from 125k to 500k. You may leverage these models to study the properties of MIM methods. Please look into the [data scaling](https://arxiv.org/abs/2206.04664) paper for more details. - -***07/09/2022*** - -`News`: - -1. SwinV2-G achieves `61.4 mIoU` on ADE20K semantic segmentation (+1.5 mIoU over the previous SwinV2-G model), using an additional [feature distillation (FD)](https://github.com/SwinTransformer/Feature-Distillation) approach, **setting a new recrod** on this benchmark. FD is an approach that can generally improve the fine-tuning performance of various pre-trained models, including DeiT, DINO, and CLIP. Particularly, it improves CLIP pre-trained ViT-L by +1.6% to reach `89.0%` on ImageNet-1K image classification, which is **the most accurate ViT-L model**. -2. Merged a PR from **Nvidia** that links to faster Swin Transformer inference that have significant speed improvements on `T4 and A100 GPUs`. -3. Merged a PR from **Nvidia** that enables an option to use `pure FP16 (Apex O2)` in training, while almost maintaining the accuracy. - -***06/03/2022*** - -1. Added **Swin-MoE**, the Mixture-of-Experts variant of Swin Transformer implemented using [Tutel](https://github.com/microsoft/tutel) (an optimized Mixture-of-Experts implementation). **Swin-MoE** is introduced in the [TuTel](https://arxiv.org/abs/2206.03382) paper. - -***05/12/2022*** - -1. Pretrained models of [Swin Transformer V2](https://arxiv.org/abs/2111.09883) on ImageNet-1K and ImageNet-22K are released. -2. ImageNet-22K pretrained models for Swin-V1-Tiny and Swin-V2-Small are released. - -***03/02/2022*** - -1. Swin Transformer V2 and SimMIM got accepted by CVPR 2022. [SimMIM](https://github.com/microsoft/SimMIM) is a self-supervised pre-training approach based on masked image modeling, a key technique that works out the 3-billion-parameter Swin V2 model using `40x less labelled data` than that of previous billion-scale models based on JFT-3B. - -***02/09/2022*** - -1. Integrated into [Huggingface Spaces 🤗](https://huggingface.co/spaces) using [Gradio](https://github.com/gradio-app/gradio). Try out the Web Demo [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/akhaliq/Swin-Transformer) - -***10/12/2021*** - -1. Swin Transformer received ICCV 2021 best paper award (Marr Prize). - -***08/09/2021*** -1. [Soft Teacher](https://arxiv.org/pdf/2106.09018v2.pdf) will appear at ICCV2021. The code will be released at [GitHub Repo](https://github.com/microsoft/SoftTeacher). `Soft Teacher` is an end-to-end semi-supervisd object detection method, achieving a new record on the COCO test-dev: `61.3 box AP` and `53.0 mask AP`. - -***07/03/2021*** -1. Add **Swin MLP**, which is an adaption of `Swin Transformer` by replacing all multi-head self-attention (MHSA) blocks by MLP layers (more precisely it is a group linear layer). The shifted window configuration can also significantly improve the performance of vanilla MLP architectures. - -***06/25/2021*** -1. [Video Swin Transformer](https://arxiv.org/abs/2106.13230) is released at [Video-Swin-Transformer](https://github.com/SwinTransformer/Video-Swin-Transformer). -`Video Swin Transformer` achieves state-of-the-art accuracy on a broad range of video recognition benchmarks, including action recognition (`84.9` top-1 accuracy on Kinetics-400 and `86.1` top-1 accuracy on Kinetics-600 with `~20x` less pre-training data and `~3x` smaller model size) and temporal modeling (`69.6` top-1 accuracy on Something-Something v2). - -***05/12/2021*** -1. Used as a backbone for `Self-Supervised Learning`: [Transformer-SSL](https://github.com/SwinTransformer/Transformer-SSL) - -Using Swin-Transformer as the backbone for self-supervised learning enables us to evaluate the transferring performance of the learnt representations on down-stream tasks, which is missing in previous works due to the use of ViT/DeiT, which has not been well tamed for down-stream tasks. - -***04/12/2021*** - -Initial commits: - -1. Pretrained models on ImageNet-1K ([Swin-T-IN1K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth), [Swin-S-IN1K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth), [Swin-B-IN1K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224.pth)) and ImageNet-22K ([Swin-B-IN22K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth), [Swin-L-IN22K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22k.pth)) are provided. -2. The supported code and models for ImageNet-1K image classification, COCO object detection and ADE20K semantic segmentation are provided. -3. The cuda kernel implementation for the [local relation layer](https://arxiv.org/pdf/1904.11491.pdf) is provided in branch [LR-Net](https://github.com/microsoft/Swin-Transformer/tree/LR-Net). - -## Introduction - -**Swin Transformer** (the name `Swin` stands for **S**hifted **win**dow) is initially described in [arxiv](https://arxiv.org/abs/2103.14030), which capably serves as a -general-purpose backbone for computer vision. It is basically a hierarchical Transformer whose representation is -computed with shifted windows. The shifted windowing scheme brings greater efficiency by limiting self-attention -computation to non-overlapping local windows while also allowing for cross-window connection. - -Swin Transformer achieves strong performance on COCO object detection (`58.7 box AP` and `51.1 mask AP` on test-dev) and -ADE20K semantic segmentation (`53.5 mIoU` on val), surpassing previous models by a large margin. - -![teaser](figures/teaser.png) - -## Main Results on ImageNet with Pretrained Models - -**ImageNet-1K and ImageNet-22K Pretrained Swin-V1 Models** - -| name | pretrain | resolution |acc@1 | acc@5 | #params | FLOPs | FPS| 22K model | 1K model | -| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |:---: |:---: | -| Swin-T | ImageNet-1K | 224x224 | 81.2 | 95.5 | 28M | 4.5G | 755 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth)/[baidu](https://pan.baidu.com/s/156nWJy4Q28rDlrX-rRbI3w)/[config](configs/swin/swin_tiny_patch4_window7_224.yaml)/[log](https://github.com/SwinTransformer/storage/files/7745562/log_swin_tiny_patch4_window7_224.txt) | -| Swin-S | ImageNet-1K | 224x224 | 83.2 | 96.2 | 50M | 8.7G | 437 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth)/[baidu](https://pan.baidu.com/s/1KFjpj3Efey3LmtE1QqPeQg)/[config](configs/swin/swin_small_patch4_window7_224.yaml)/[log](https://github.com/SwinTransformer/storage/files/7745563/log_swin_small_patch4_window7_224.txt) | -| Swin-B | ImageNet-1K | 224x224 | 83.5 | 96.5 | 88M | 15.4G | 278 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224.pth)/[baidu](https://pan.baidu.com/s/16bqCTEc70nC_isSsgBSaqQ)/[config](configs/swin/swin_base_patch4_window7_224.yaml)/[log](https://github.com/SwinTransformer/storage/files/7745564/log_swin_base_patch4_window7_224.txt) | -| Swin-B | ImageNet-1K | 384x384 | 84.5 | 97.0 | 88M | 47.1G | 85 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384.pth)/[baidu](https://pan.baidu.com/s/1xT1cu740-ejW7htUdVLnmw)/[config](configs/swin/swin_base_patch4_window12_384_finetune.yaml) | -| Swin-T | ImageNet-22K | 224x224 | 80.9 | 96.0 | 28M | 4.5G | 755 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.8/swin_tiny_patch4_window7_224_22k.pth)/[baidu](https://pan.baidu.com/s/1vct0VYwwQQ8PYkBjwSSBZQ?pwd=swin)/[config](configs/swin/swin_tiny_patch4_window7_224_22k.yaml) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.8/swin_tiny_patch4_window7_224_22kto1k_finetune.pth)/[baidu](https://pan.baidu.com/s/1K0OO-nGZDPkR8fm_r83e8Q?pwd=swin)/[config](configs/swin/swin_tiny_patch4_window7_224_22kto1k_finetune.yaml) | -| Swin-S | ImageNet-22K | 224x224 | 83.2 | 97.0 | 50M | 8.7G | 437 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.8/swin_small_patch4_window7_224_22k.pth)/[baidu](https://pan.baidu.com/s/11NC1xdT5BAGBgazdTme5Sg?pwd=swin)/[config](configs/swin/swin_small_patch4_window7_224_22k.yaml) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.8/swin_small_patch4_window7_224_22kto1k_finetune.pth)/[baidu](https://pan.baidu.com/s/10RFVfjQJhwPfeHrmxQUaLw?pwd=swin)/[config](configs/swin/swin_small_patch4_window7_224_22kto1k_finetune.yaml) | -| Swin-B | ImageNet-22K | 224x224 | 85.2 | 97.5 | 88M | 15.4G | 278 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth)/[baidu](https://pan.baidu.com/s/1y1Ec3UlrKSI8IMtEs-oBXA)/[config](configs/swin/swin_base_patch4_window7_224_22k.yaml) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22kto1k.pth)/[baidu](https://pan.baidu.com/s/1n_wNkcbRxVXit8r_KrfAVg)/[config](configs/swin/swin_base_patch4_window7_224_22kto1k_finetune.yaml) | -| Swin-B | ImageNet-22K | 384x384 | 86.4 | 98.0 | 88M | 47.1G | 85 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22k.pth)/[baidu](https://pan.baidu.com/s/1vwJxnJcVqcLZAw9HaqiR6g) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22kto1k.pth)/[baidu](https://pan.baidu.com/s/1caKTSdoLJYoi4WBcnmWuWg)/[config](configs/swin/swin_base_patch4_window12_384_22kto1k_finetune.yaml) | -| Swin-L | ImageNet-22K | 224x224 | 86.3 | 97.9 | 197M | 34.5G | 141 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22k.pth)/[baidu](https://pan.baidu.com/s/1pws3rOTFuOebBYP3h6Kx8w)/[config](configs/swin/swin_large_patch4_window7_224_22k.yaml) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22kto1k.pth)/[baidu](https://pan.baidu.com/s/1NkQApMWUhxBGjk1ne6VqBQ)/[config](configs/swin/swin_large_patch4_window7_224_22kto1k_finetune.yaml) | -| Swin-L | ImageNet-22K | 384x384 | 87.3 | 98.2 | 197M | 103.9G | 42 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22k.pth)/[baidu](https://pan.baidu.com/s/1sl7o_bJA143OD7UqSLAMoA) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22kto1k.pth)/[baidu](https://pan.baidu.com/s/1X0FLHQyPOC6Kmv2CmgxJvA)/[config](configs/swin/swin_large_patch4_window12_384_22kto1k_finetune.yaml) | - -**ImageNet-1K and ImageNet-22K Pretrained Swin-V2 Models** - -| name | pretrain | resolution | window |acc@1 | acc@5 | #params | FLOPs | FPS |22K model | 1K model | -|:---------------------:| :---: | :---: | :---: | :---: | :---: | :---: | :---: |:---:|:---: |:---: | -| SwinV2-T | ImageNet-1K | 256x256 | 8x8 | 81.8 | 95.9 | 28M | 5.9G | 572 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_tiny_patch4_window8_256.pth)/[baidu](https://pan.baidu.com/s/1RzLkAH_5OtfRCJe6Vlg6rg?pwd=swin)/[config](configs/swinv2/swinv2_tiny_patch4_window8_256.yaml) | -| SwinV2-S | ImageNet-1K | 256x256 | 8x8 | 83.7 | 96.6 | 50M | 11.5G | 327 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_small_patch4_window8_256.pth)/[baidu](https://pan.baidu.com/s/195PdA41szEduW3jEtRSa4Q?pwd=swin)/[config](configs/swinv2/swinv2_small_patch4_window8_256.yaml) | -| SwinV2-B | ImageNet-1K | 256x256 | 8x8 | 84.2 | 96.9 | 88M | 20.3G | 217 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_base_patch4_window8_256.pth)/[baidu](https://pan.baidu.com/s/18AfMSz3dPyzIvP1dKuERvQ?pwd=swin)/[config](configs/swinv2/swinv2_base_patch4_window8_256.yaml) | -| SwinV2-T | ImageNet-1K | 256x256 | 16x16 | 82.8 | 96.2 | 28M | 6.6G | 437 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_tiny_patch4_window16_256.pth)/[baidu](https://pan.baidu.com/s/1dyK3cK9Xipmv6RnTtrPocw?pwd=swin)/[config](configs/swinv2/swinv2_tiny_patch4_window16_256.yaml) | -| SwinV2-S | ImageNet-1K | 256x256 | 16x16 | 84.1 | 96.8 | 50M | 12.6G | 257 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_small_patch4_window16_256.pth)/[baidu](https://pan.baidu.com/s/1ZIPiSfWNKTPp821Ka-Mifw?pwd=swin)/[config](configs/swinv2/swinv2_small_patch4_window16_256.yaml) | -| SwinV2-B | ImageNet-1K | 256x256 | 16x16 | 84.6 | 97.0 | 88M | 21.8G | 174 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_base_patch4_window16_256.pth)/[baidu](https://pan.baidu.com/s/1dlDQGn8BXCmnh7wQSM5Nhw?pwd=swin)/[config](configs/swinv2/swinv2_base_patch4_window16_256.yaml) | -| SwinV2-B\* | ImageNet-22K | 256x256 | 16x16 | 86.2 | 97.9 | 88M | 21.8G | 174 | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_base_patch4_window12_192_22k.pth)/[baidu](https://pan.baidu.com/s/1Xc2rsSsRQz_sy5mjgfxrMQ?pwd=swin)/[config](configs/swinv2/swinv2_base_patch4_window12_192_22k.yaml) | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_base_patch4_window12to16_192to256_22kto1k_ft.pth)/[baidu](https://pan.baidu.com/s/1sgstld4MgGsZxhUAW7MlmQ?pwd=swin)/[config](configs/swinv2/swinv2_base_patch4_window12to16_192to256_22kto1k_ft.yaml) | -| SwinV2-B\* | ImageNet-22K | 384x384 | 24x24 | 87.1 | 98.2 | 88M | 54.7G | 57 | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_base_patch4_window12_192_22k.pth)/[baidu](https://pan.baidu.com/s/1Xc2rsSsRQz_sy5mjgfxrMQ?pwd=swin)/[config](configs/swinv2/swinv2_base_patch4_window12_192_22k.yaml) | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_base_patch4_window12to24_192to384_22kto1k_ft.pth)/[baidu](https://pan.baidu.com/s/17u3sEQaUYlvfL195rrORzQ?pwd=swin)/[config](configs/swinv2/swinv2_base_patch4_window12to24_192to384_22kto1k_ft.yaml) | -| SwinV2-L\* | ImageNet-22K | 256x256 | 16x16 | 86.9 | 98.0 | 197M | 47.5G | 95 | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_large_patch4_window12_192_22k.pth)/[baidu](https://pan.baidu.com/s/11PhCV7qAGXtZ8dXNgyiGOw?pwd=swin)/[config](configs/swinv2/swinv2_large_patch4_window12_192_22k.yaml) | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_large_patch4_window12to16_192to256_22kto1k_ft.pth)/[baidu](https://pan.baidu.com/s/1pqp31N80qIWjFPbudzB6Bw?pwd=swin)/[config](configs/swinv2/swinv2_large_patch4_window12to16_192to256_22kto1k_ft.yaml) | -| SwinV2-L\* | ImageNet-22K | 384x384 | 24x24 | 87.6 | 98.3 | 197M | 115.4G | 33 | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_large_patch4_window12_192_22k.pth)/[baidu](https://pan.baidu.com/s/11PhCV7qAGXtZ8dXNgyiGOw?pwd=swin)/[config](configs/swinv2/swinv2_large_patch4_window12_192_22k.yaml) | [github](https://github.com/SwinTransformer/storage/releases/download/v2.0.0/swinv2_large_patch4_window12to24_192to384_22kto1k_ft.pth)/[baidu](https://pan.baidu.com/s/13URdNkygr3Xn0N3e6IwjgA?pwd=swin)/[config](configs/swinv2/swinv2_large_patch4_window12to24_192to384_22kto1k_ft.yaml) | - -Note: -- SwinV2-B\* (SwinV2-L\*) with input resolution of 256x256 and 384x384 both fine-tuned from the same pre-training model using a smaller input resolution of 192x192. -- SwinV2-B\* (384x384) achieves 78.08 acc@1 on ImageNet-1K-V2 while SwinV2-L\* (384x384) achieves 78.31. - -**ImageNet-1K Pretrained Swin MLP Models** +```sh +python -m venv venv +# Activate your virtual environment somehow +source venv/bin/activate.fish +``` -| name | pretrain | resolution |acc@1 | acc@5 | #params | FLOPs | FPS | 1K model | -| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | -| [Mixer-B/16](https://arxiv.org/pdf/2105.01601.pdf) | ImageNet-1K | 224x224 | 76.4 | - | 59M | 12.7G | - | [official repo](https://github.com/google-research/vision_transformer) | -| [ResMLP-S24](https://arxiv.org/abs/2105.03404) | ImageNet-1K | 224x224 | 79.4 | - | 30M | 6.0G | 715 | [timm](https://github.com/rwightman/pytorch-image-models) | -| [ResMLP-B24](https://arxiv.org/abs/2105.03404) | ImageNet-1K | 224x224 | 81.0 | - | 116M | 23.0G | 231 | [timm](https://github.com/rwightman/pytorch-image-models) | -| Swin-T/C24 | ImageNet-1K | 256x256 | 81.6 | 95.7 | 28M | 5.9G | 563 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.5/swin_tiny_c24_patch4_window8_256.pth)/[baidu](https://pan.baidu.com/s/17k-7l6Sxt7uZ7IV0f26GNQ)/[config](configs/swin/swin_tiny_c24_patch4_window8_256.yaml) | -| SwinMLP-T/C24 | ImageNet-1K | 256x256 | 79.4 | 94.6 | 20M | 4.0G | 807 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.5/swin_mlp_tiny_c24_patch4_window8_256.pth)/[baidu](https://pan.baidu.com/s/1Sa4vP5R0M2RjfIe9HIga-Q)/[config](configs/swin/swin_mlp_tiny_c24_patch4_window8_256.yaml) | -| SwinMLP-T/C12 | ImageNet-1K | 256x256 | 79.6 | 94.7 | 21M | 4.0G | 792 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.5/swin_mlp_tiny_c12_patch4_window8_256.pth)/[baidu](https://pan.baidu.com/s/1mM9J2_DEVZHUB5ASIpFl0w)/[config](configs/swin/swin_mlp_tiny_c12_patch4_window8_256.yaml) | -| SwinMLP-T/C6 | ImageNet-1K | 256x256 | 79.7 | 94.9 | 23M | 4.0G | 766 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.5/swin_mlp_tiny_c6_patch4_window8_256.pth)/[baidu](https://pan.baidu.com/s/1hUTYVT2W1CsjICw-3W-Vjg)/[config](configs/swin/swin_mlp_tiny_c6_patch4_window8_256.yaml) | -| SwinMLP-B | ImageNet-1K | 224x224 | 81.3 | 95.3 | 61M | 10.4G | 409 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.5/swin_mlp_base_patch4_window7_224.pth)/[baidu](https://pan.baidu.com/s/1zww3dnbX3GxNiGfb-GwyUg)/[config](configs/swin/swin_mlp_base_patch4_window7_224.yaml) | +CUDA 11.6 -Note: access code for `baidu` is `swin`. C24 means each head has 24 channels. +```sh +pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116 +``` -**ImageNet-22K Pretrained Swin-MoE Models** +CUDA 11.3 -- Please refer to [get_started](get_started.md#mixture-of-experts-support) for instructions on running Swin-MoE. -- Pretrained models for Swin-MoE can be found in [MODEL HUB](MODELHUB.md#imagenet-22k-pretrained-swin-moe-models) +```sh +pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113 +``` -## Main Results on Downstream Tasks +Python packages -**COCO Object Detection (2017 val)** +```sh +pip install matplotlib yacs timm einops black isort flake8 flake8-bugbear termcolor tensorboard preface opencv-python +``` -| Backbone | Method | pretrain | Lr Schd | box mAP | mask mAP | #params | FLOPs | -| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | -| Swin-T | Mask R-CNN | ImageNet-1K | 3x | 46.0 | 41.6 | 48M | 267G | -| Swin-S | Mask R-CNN | ImageNet-1K | 3x | 48.5 | 43.3 | 69M | 359G | -| Swin-T | Cascade Mask R-CNN | ImageNet-1K | 3x | 50.4 | 43.7 | 86M | 745G | -| Swin-S | Cascade Mask R-CNN | ImageNet-1K | 3x | 51.9 | 45.0 | 107M | 838G | -| Swin-B | Cascade Mask R-CNN | ImageNet-1K | 3x | 51.9 | 45.0 | 145M | 982G | -| Swin-T | RepPoints V2 | ImageNet-1K | 3x | 50.0 | - | 45M | 283G | -| Swin-T | Mask RepPoints V2 | ImageNet-1K | 3x | 50.3 | 43.6 | 47M | 292G | -| Swin-B | HTC++ | ImageNet-22K | 6x | 56.4 | 49.1 | 160M | 1043G | -| Swin-L | HTC++ | ImageNet-22K | 3x | 57.1 | 49.5 | 284M | 1470G | -| Swin-L | HTC++* | ImageNet-22K | 3x | 58.0 | 50.4 | 284M | - | +2. Install Apex -Note: * indicates multi-scale testing. +```sh +git clone https://github.com/NVIDIA/apex.git +cd apex +pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./ +``` -**ADE20K Semantic Segmentation (val)** +```sh +cd kernels/window_process +python setup.py install +``` -| Backbone | Method | pretrain | Crop Size | Lr Schd | mIoU | mIoU (ms+flip) | #params | FLOPs | -| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | -| Swin-T | UPerNet | ImageNet-1K | 512x512 | 160K | 44.51 | 45.81 | 60M | 945G | -| Swin-S | UperNet | ImageNet-1K | 512x512 | 160K | 47.64 | 49.47 | 81M | 1038G | -| Swin-B | UperNet | ImageNet-1K | 512x512 | 160K | 48.13 | 49.72 | 121M | 1188G | -| Swin-B | UPerNet | ImageNet-22K | 640x640 | 160K | 50.04 | 51.66 | 121M | 1841G | -| Swin-L | UperNet | ImageNet-22K | 640x640 | 160K | 52.05 | 53.53 | 234M | 3230G | +3. Download Data -## Citing Swin Transformer +We use the iNat21 dataseta available on [GitHub](https://github.com/visipedia/inat_comp/tree/master/2021) ``` -@inproceedings{liu2021Swin, - title={Swin Transformer: Hierarchical Vision Transformer using Shifted Windows}, - author={Liu, Ze and Lin, Yutong and Cao, Yue and Hu, Han and Wei, Yixuan and Zhang, Zheng and Lin, Stephen and Guo, Baining}, - booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)}, - year={2021} -} -``` -## Citing Local Relation Networks (the first full-attention visual backbone) +mkdir -p data/inat21 +cd data/inat21 +mkdir compressed raw +cd compressed +wget https://ml-inat-competition-datasets.s3.amazonaws.com/2021/train.tar.gz +wget https://ml-inat-competition-datasets.s3.amazonaws.com/2021/val.tar.gz + +# pv is just a progress bar +pv val.tar.gz | tar -xz +mv val ../raw/ # if I knew how tar worked I could have it extract to raw/ + +pv train.tar.gz | tar -xz +mv train ../raw/ ``` -@inproceedings{hu2019local, - title={Local Relation Networks for Image Recognition}, - author={Hu, Han and Zhang, Zheng and Xie, Zhenda and Lin, Stephen}, - booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)}, - pages={3464--3473}, - year={2019} -} -``` -## Citing Swin Transformer V2 -``` -@inproceedings{liu2021swinv2, - title={Swin Transformer V2: Scaling Up Capacity and Resolution}, - author={Ze Liu and Han Hu and Yutong Lin and Zhuliang Yao and Zhenda Xie and Yixuan Wei and Jia Ning and Yue Cao and Zheng Zhang and Li Dong and Furu Wei and Baining Guo}, - booktitle={International Conference on Computer Vision and Pattern Recognition (CVPR)}, - year={2022} -} -``` -## Citing SimMIM (a self-supervised approach that enables SwinV2-G) -``` -@inproceedings{xie2021simmim, - title={SimMIM: A Simple Framework for Masked Image Modeling}, - author={Xie, Zhenda and Zhang, Zheng and Cao, Yue and Lin, Yutong and Bao, Jianmin and Yao, Zhuliang and Dai, Qi and Hu, Han}, - booktitle={International Conference on Computer Vision and Pattern Recognition (CVPR)}, - year={2022} -} -``` -## Citing SimMIM-data-scaling -``` -@article{xie2022data, - title={On Data Scaling in Masked Image Modeling}, - author={Xie, Zhenda and Zhang, Zheng and Cao, Yue and Lin, Yutong and Wei, Yixuan and Dai, Qi and Hu, Han}, - journal={arXiv preprint arXiv:2206.04664}, - year={2022} -} -``` -## Citing Swin-MoE -``` -@misc{hwang2022tutel, - title={Tutel: Adaptive Mixture-of-Experts at Scale}, - author={Changho Hwang and Wei Cui and Yifan Xiong and Ziyue Yang and Ze Liu and Han Hu and Zilong Wang and Rafael Salas and Jithin Jose and Prabhat Ram and Joe Chau and Peng Cheng and Fan Yang and Mao Yang and Yongqiang Xiong}, - year={2022}, - eprint={2206.03382}, - archivePrefix={arXiv} -} -``` - -## Getting Started -- For **Image Classification**, please see [get_started.md](get_started.md) for detailed instructions. -- For **Object Detection and Instance Segmentation**, please see [Swin Transformer for Object Detection](https://github.com/SwinTransformer/Swin-Transformer-Object-Detection). -- For **Semantic Segmentation**, please see [Swin Transformer for Semantic Segmentation](https://github.com/SwinTransformer/Swin-Transformer-Semantic-Segmentation). -- For **Self-Supervised Learning**, please see [Transformer-SSL](https://github.com/SwinTransformer/Transformer-SSL). -- For **Video Recognition**, please see [Video Swin Transformer](https://github.com/SwinTransformer/Video-Swin-Transformer). +4. Preprocess iNat 21 -## Third-party Usage and Experiments +Use your root data folder and your size of choice. -***In this pargraph, we cross link third-party repositories which use Swin and report results. You can let us know by raising an issue*** - -(`Note please report accuracy numbers and provide trained models in your new repository to facilitate others to get sense of correctness and model behavior`) - -[06/30/2022] Swin Transformers (V1) inference implemented in FasterTransformer: [FasterTransformer](https://github.com/NVIDIA/FasterTransformer/blob/main/docs/swin_guide.md) - -[05/12/2022] Swin Transformers (V1) implemented in TensorFlow with the pre-trained parameters ported into them. Find the implementation, -TensorFlow weights, code example here in [this repository](https://github.com/sayakpaul/swin-transformers-tf/). - -[04/06/2022] Swin Transformer for Audio Classification: [Hierarchical Token Semantic Audio Transformer](https://github.com/RetroCirce/HTS-Audio-Transformer). - -[12/21/2021] Swin Transformer for StyleGAN: [StyleSwin](https://github.com/microsoft/StyleSwin) - -[12/13/2021] Swin Transformer for Face Recognition: [FaceX-Zoo](https://github.com/JDAI-CV/FaceX-Zoo) - -[08/29/2021] Swin Transformer for Image Restoration: [SwinIR](https://github.com/JingyunLiang/SwinIR) - -[08/12/2021] Swin Transformer for person reID: [https://github.com/layumi/Person_reID_baseline_pytorch](https://github.com/layumi/Person_reID_baseline_pytorch) - -[06/29/2021] Swin-Transformer in PaddleClas and inference based on whl package: [https://github.com/PaddlePaddle/PaddleClas](https://github.com/PaddlePaddle/PaddleClas) - -[04/14/2021] Swin for RetinaNet in Detectron: https://github.com/xiaohu2015/SwinT_detectron2. - -[04/16/2021] Included in a famous model zoo: https://github.com/rwightman/pytorch-image-models. - -[04/20/2021] Swin-Transformer classifier inference using TorchServe: https://github.com/kamalkraj/Swin-Transformer-Serve - -## Contributing - -This project welcomes contributions and suggestions. Most contributions require you to agree to a -Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us -the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com. - -When you submit a pull request, a CLA bot will automatically determine whether you need to provide -a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions -provided by the bot. You will only need to do this once across all repos using our CLA. - -This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). -For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or -contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments. - -## Trademarks - -This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft -trademarks or logos is subject to and must follow -[Microsoft's Trademark & Brand Guidelines](https://www.microsoft.com/en-us/legal/intellectualproperty/trademarks/usage/general). -Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship. -Any use of third-party trademarks or logos are subject to those third-party's policies. +``` +python -m data.inat preprocess /mnt/10tb/data/inat21/ val resize 224 +python -m data.inat preprocess /mnt/10tb/data/inat21/ train resize 224 +``` From 13fcf53bea4d5f8601dc0912497c7f01eadc57ec Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Thu, 27 Oct 2022 22:49:14 +0000 Subject: [PATCH 13/47] Add inat config --- ...ase_patch4_window12_192_inat21_lr1.25.yaml | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 configs/swinv2/swinv2_base_patch4_window12_192_inat21_lr1.25.yaml diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_lr1.25.yaml b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_lr1.25.yaml new file mode 100644 index 00000000..516be54e --- /dev/null +++ b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_lr1.25.yaml @@ -0,0 +1,33 @@ +DATA: + DATASET: inat21 + IMG_SIZE: 192 + DATA_PATH: /mnt/10tb/data/inat21/resize-192 + NUM_WORKERS: 32 + BATCH_SIZE: 64 +MODEL: + TYPE: swinv2 + NAME: swinv2_base_patch4_window12_192_inat21 + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 128 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 4, 8, 16, 32 ] + WINDOW_SIZE: 12 +TRAIN: + # Want a global batch size of 2048 because SwinV2 was trained on 16 V100s with batch size 128 (I think) + # But we are going to use a global batch size of 1024 because it's faster (throughput). + ACCUMULATION_STEPS: 2 + + # We are using limited epochs based on pre-training configs for imagenet22k + # Then we will pre-train on 256x256 for 30 epochs + EPOCHS: 90 + WARMUP_EPOCHS: 5 + WEIGHT_DECAY: 0.1 + + # Use 1/4 of original learning rates + BASE_LR: 1.25e-4 + WARMUP_LR: 1.25e-7 + MIN_LR: 1.25e-6 + HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] + +HIERARCHICAL: true From 9880bb55d64b61c29cf91cf1907718e56ab587d9 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Thu, 27 Oct 2022 22:49:21 +0000 Subject: [PATCH 14/47] Fix bad logging --- logger.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/logger.py b/logger.py index 49d44722..952b2898 100644 --- a/logger.py +++ b/logger.py @@ -58,13 +58,10 @@ def __init__(self, output_dir, dist_rank): self.output_dir = output_dir if dist_rank == 0: - self.writer = CorrectedSummaryWriter(log_dir=self.output_dir) + self.writer = SummaryWriter(log_dir=self.output_dir) def add_hparams(self, hparams, metrics): - if self.writer is None: - return - - self.writer.add_hparams(hparams, metrics, run_name="") + pass def log(self, items, step): if self.writer is None: From 2a97f6b98f1af23d97bc410f3827964a4bce6fe2 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Sat, 29 Oct 2022 21:51:17 +0000 Subject: [PATCH 15/47] Add hierarchical lr 2.5 config --- ...indow12_192_inat21_hierarchical_lr2.5.yaml | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml new file mode 100644 index 00000000..9deae81e --- /dev/null +++ b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml @@ -0,0 +1,33 @@ +DATA: + DATASET: inat21 + IMG_SIZE: 192 + DATA_PATH: /mnt/10tb/data/inat21/resize-192 + NUM_WORKERS: 32 + BATCH_SIZE: 64 +MODEL: + TYPE: swinv2 + NAME: swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5 + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 128 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 4, 8, 16, 32 ] + WINDOW_SIZE: 12 +TRAIN: + # Want a global batch size of 2048 because SwinV2 was trained on 16 V100s with batch size 128 (I think) + # But we are going to use a global batch size of 1024 because it's faster (throughput). + ACCUMULATION_STEPS: 2 + + # We are using limited epochs based on pre-training configs for imagenet22k + # Then we will pre-train on 256x256 for 30 epochs + EPOCHS: 90 + WARMUP_EPOCHS: 5 + WEIGHT_DECAY: 0.1 + + # Use 1/2 of original learning rates + BASE_LR: 2.5e-4 + WARMUP_LR: 2.5e-7 + MIN_LR: 2.5e-6 + HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] + +HIERARHICAL: true From d0e7415c895a3c8fc296e6416b367deea98b3342 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Sat, 29 Oct 2022 21:51:32 +0000 Subject: [PATCH 16/47] Fix constants for inat 192/256 --- data/constants.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/data/constants.py b/data/constants.py index 3919bdc4..c1e1ef5e 100644 --- a/data/constants.py +++ b/data/constants.py @@ -2,11 +2,15 @@ data_mean_std = { "/mnt/10tb/data/inat21/resize-192": ( - torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), + torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), ), "/mnt/10tb/data/inat21/resize-224": ( torch.tensor([0.23762744665145874, 0.2292044311761856, 0.24757201969623566]), torch.tensor([0.4632636606693268, 0.48004215955734253, 0.37622377276420593]), ), + "/mnt/10tb/data/inat21/resize-256": ( + torch.tensor([0.23768986761569977, 0.22925858199596405, 0.2476460039615631]), + torch.tensor([0.4632672071456909, 0.480050653219223, 0.37618669867515564]), + ), } From fdf6d331c2b446be4a1350c196595ec84ce5d88f Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sat, 29 Oct 2022 22:10:10 +0000 Subject: [PATCH 17/47] rename experiment --- ...v2_base_patch4_window12_192_inat21_hierarchical_lr1.25.yaml} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename configs/swinv2/{swinv2_base_patch4_window12_192_inat21_lr1.25.yaml => swinv2_base_patch4_window12_192_inat21_hierarchical_lr1.25.yaml} (92%) diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_lr1.25.yaml b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr1.25.yaml similarity index 92% rename from configs/swinv2/swinv2_base_patch4_window12_192_inat21_lr1.25.yaml rename to configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr1.25.yaml index 516be54e..cc23379a 100644 --- a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_lr1.25.yaml +++ b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr1.25.yaml @@ -6,7 +6,7 @@ DATA: BATCH_SIZE: 64 MODEL: TYPE: swinv2 - NAME: swinv2_base_patch4_window12_192_inat21 + NAME: swinv2_base_patch4_window12_192_inat21_hierarchical_lr1.25 DROP_PATH_RATE: 0.2 SWINV2: EMBED_DIM: 128 From 424f78bfb7f0d2a267c16c93f8115aea6ff85ee4 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 18:30:04 +0000 Subject: [PATCH 18/47] Move imagenet mean/std to constants.py --- data/build.py | 1 - data/constants.py | 2 ++ utils.py | 24 ------------------------ 3 files changed, 2 insertions(+), 25 deletions(-) diff --git a/data/build.py b/data/build.py index e261f454..c9b140cd 100644 --- a/data/build.py +++ b/data/build.py @@ -12,7 +12,6 @@ import torch import torch.distributed as dist from timm.data import Mixup, create_transform -from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD from torch.utils.data import Subset from torchvision import datasets, transforms diff --git a/data/constants.py b/data/constants.py index c1e1ef5e..5bb8b2eb 100644 --- a/data/constants.py +++ b/data/constants.py @@ -1,4 +1,5 @@ import torch +from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD data_mean_std = { "/mnt/10tb/data/inat21/resize-192": ( @@ -13,4 +14,5 @@ torch.tensor([0.23768986761569977, 0.22925858199596405, 0.2476460039615631]), torch.tensor([0.4632672071456909, 0.480050653219223, 0.37618669867515564]), ), + "imagenet": (IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD), } diff --git a/utils.py b/utils.py index 29c664d8..b4b7c79e 100644 --- a/utils.py +++ b/utils.py @@ -7,7 +7,6 @@ import os -import preface.dict import torch import torch.distributed as dist from torch._six import inf @@ -289,26 +288,3 @@ def state_dict(self): def load_state_dict(self, state_dict): self._scaler.load_state_dict(state_dict) - - -def to_hparams(cfg): - hparams = {} - for key, value in cfg.items(): - key = key.lower() - if ( - isinstance(value, int) - or isinstance(value, float) - or isinstance(value, str) - or isinstance(value, bool) - or isinstance(value, torch.Tensor) - or value is None - ): - hparams[key] = value - elif isinstance(value, dict): - hparams[key] = to_hparams(value) - elif isinstance(value, list) or isinstance(value, tuple): - hparams[key] = "[" + ", ".join(repr(v) for v in value) + "]" - else: - raise ValueError(f"Don't know how to handle {key}: {value}!") - - return preface.dict.flattened(hparams) From 0d326942b1f42a3d0bc2a28c49f9781ae5122aa8 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 18:31:06 +0000 Subject: [PATCH 19/47] Add wandb/remove tensorboard --- logger.py | 40 ++++++++++++++++++++++++++++------------ main.py | 18 ++++++++++-------- 2 files changed, 38 insertions(+), 20 deletions(-) diff --git a/logger.py b/logger.py index 952b2898..a557d2d3 100644 --- a/logger.py +++ b/logger.py @@ -11,9 +11,9 @@ import sys import torch +import wandb from termcolor import colored from torch.utils.tensorboard import SummaryWriter -from torch.utils.tensorboard.summary import hparams @functools.lru_cache() @@ -82,15 +82,31 @@ def log(self, items, step): print(f"Can't log {v} because it is {type(v)}!") -class CorrectedSummaryWriter(SummaryWriter): - def add_hparams(self, hparam_dict, metric_dict): - torch._C._log_api_usage_once("tensorboard.logging.add_hparams") - if type(hparam_dict) is not dict or type(metric_dict) is not dict: - raise TypeError("hparam_dict and metric_dict should be dictionary.") - exp, ssi, sei = hparams(hparam_dict, metric_dict) +class WandbWriter: + def __init__(self, rank): + self.rank = rank - self.file_writer.add_summary(exp) - self.file_writer.add_summary(ssi) - self.file_writer.add_summary(sei) - for k, v in metric_dict.items(): - self.add_scalar(k, v) + def init(self, config): + if self.rank != 0: + return + + wandb.init( + config=config, project="hierarchical-vision", dir="./runs", resume=True + ) + + wandb.define_metric("val/loss", step_metric="epoch", summary="min") + wandb.define_metric("val/acc1", step_metric="epoch", summary="max") + wandb.define_metric("val/acc5", step_metric="epoch", summary="max") + + def log(self, dct): + if self.rank != 0: + return + + wandb.log(dct) + + @property + def name(self): + if self.rank != 0: + raise RuntimeError(f"Should not get .name with rank {self.rank}.") + + return wandb.run.name diff --git a/main.py b/main.py index 98ac6a3e..3aa0d582 100644 --- a/main.py +++ b/main.py @@ -26,7 +26,7 @@ HierarchicalCrossEntropyLoss, accuracy, ) -from logger import TensorboardWriter, create_logger +from logger import WandbWriter, create_logger from lr_scheduler import build_scheduler from models import build_model from optimizer import build_optimizer @@ -38,7 +38,6 @@ load_pretrained, reduce_tensor, save_checkpoint, - to_hparams, ) @@ -374,14 +373,15 @@ def train_one_epoch( // config.TRAIN.ACCUMULATION_STEPS )[0] - tb_writer.log(stats, (epoch * num_steps + idx)) + wandb_writer.log( + {**stats, "step": epoch * num_steps + idx, "epoch": epoch, "batch": idx} + ) epoch_time = time.time() - start logger.info( f"EPOCH {epoch} training took {datetime.timedelta(seconds=int(epoch_time))}" ) - tb_writer.log({"train_time": epoch_time}, epoch) - tb_writer.add_hparams(to_hparams(config), stats) + wandb_writer.log({"train_time": epoch_time, "epoch": epoch}, epoch) @torch.no_grad() @@ -433,14 +433,15 @@ def validate(config, data_loader, model, epoch): f"Mem {memory_used:.0f}MB" ) logger.info(f" * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}") - tb_writer.log( + wandb_writer.log( { "val_acc1": acc1_meter.avg, "val_acc5": acc5_meter.avg, "val_loss": loss_meter.avg, + "epoch": epoch, }, - epoch, ) + return acc1_meter.avg, acc5_meter.avg, loss_meter.avg @@ -519,7 +520,8 @@ def throughput(data_loader, model, logger): logger = create_logger( output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" ) - tb_writer = TensorboardWriter(output_dir=config.OUTPUT, dist_rank=dist.get_rank()) + wandb_writer = WandbWriter(rank=dist.get_rank()) + wandb_writer.init(config) if dist.get_rank() == 0: path = os.path.join(config.OUTPUT, "config.yaml") From fda74432313b759534cf30c4bb4d7b689e3b866a Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 18:31:24 +0000 Subject: [PATCH 20/47] Fix lints --- main.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/main.py b/main.py index 3aa0d582..e562778d 100644 --- a/main.py +++ b/main.py @@ -120,7 +120,7 @@ def parse_option(): parser.add_argument( "--fused_layernorm", action="store_true", help="Use fused layernorm." ) - ## overwrite optimizer in config (*.yaml) if specified, e.g., fused_adam/fused_lamb + # overwrite optimizer in config (*.yaml) if specified, e.g., fused_adam/fused_lamb parser.add_argument( "--optim", type=str, @@ -452,12 +452,12 @@ def throughput(data_loader, model, logger): for idx, (images, _) in enumerate(data_loader): images = images.cuda(non_blocking=True) batch_size = images.shape[0] - for i in range(50): + for _ in range(50): model(images) torch.cuda.synchronize() - logger.info(f"throughput averaged with 30 times") + logger.info("throughput averaged with 30 times") tic1 = time.time() - for i in range(30): + for _ in range(30): model(images) torch.cuda.synchronize() tic2 = time.time() From 86e21c92fffb104b034641a22219752c32c13fcd Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 18:33:55 +0000 Subject: [PATCH 21/47] Make it possible to start from hierarchical checkpoints --- utils.py | 178 ++++++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 157 insertions(+), 21 deletions(-) diff --git a/utils.py b/utils.py index b4b7c79e..f7a756aa 100644 --- a/utils.py +++ b/utils.py @@ -49,6 +49,156 @@ def load_checkpoint(config, model, optimizer, lr_scheduler, loss_scaler, logger) return max_accuracy +def _hierarchical_weight_k(i): + return f"head.heads.{i}.weight" + + +def _hierarchical_bias_k(i): + return f"head.heads.{i}.bias" + + +def handle_linear_head(config, model, state_dict, logger): + """ + Check classifier, if not match, then re-init classifier to zero + Ways it could not match: + + 1. Both have a single linear head, with different number of classes + + 2. Pretrained with hierarchical head and are now finetuning with single linear head. + If the fine-grained head was 10K classes and the linear head has 10K classes, + and the current dataset is iNat, we use the pre-trained fine-grained head. + + 3. Pretrained with single linear head and are finetuning with hierarchical head + (for instance, if starting with imagenet pre-training, then doing domain-specific + pre-training on iNat21) + + 4. Both have a hierarchical head with different number of tiers/classes. + We always reinitialize the hierarchical head. + """ + + pretrained_hierarchical = _hierarchical_bias_k(0) in state_dict + current_hierarchical = config.HIERARCHICAL + + if not pretrained_hierarchical and not current_hierarchical: + # TESTED because Microsoft wrote this code. + # Both have a single linear head + assert "head.bias" in state_dict, "Should have a single pre-trained linear head" + assert hasattr(model.head, "bias"), "Should have a single random linear head" + + head_bias_pretrained = state_dict["head.bias"] + num_classes_pretrained = head_bias_pretrained.shape[0] + num_classes = model.head.bias.shape[0] + if num_classes_pretrained != num_classes: + if num_classes_pretrained == 21841 and num_classes == 1000: + logger.info("loading ImageNet-22K weight to ImageNet-1K ......") + map22kto1k_path = "data/map22kto1k.txt" + with open(map22kto1k_path) as f: + map22kto1k = f.readlines() + map22kto1k = [int(id22k.strip()) for id22k in map22kto1k] + state_dict["head.weight"] = state_dict["head.weight"][map22kto1k, :] + state_dict["head.bias"] = state_dict["head.bias"][map22kto1k] + else: + torch.nn.init.constant_(model.head.bias, 0.0) + torch.nn.init.constant_(model.head.weight, 0.0) + del state_dict["head.weight"] + del state_dict["head.bias"] + logger.warning( + "Error in loading classifier head, re-init classifier head to 0" + ) + elif pretrained_hierarchical and not current_hierarchical: + # UNTESTED + assert ( + "head.bias" not in state_dict + ), "Should not have a single pre-trained linear head" + assert hasattr(model.head, "bias"), "Should have a single random linear head" + # Increment finegrained level until the key doesn't exist. + # Then it is the last level in the hierarchical model + max_level = 0 + while _hierarchical_bias_k(max_level) in state_dict: + max_level += 1 + + finegrained_num_classes_pretrained = state_dict[ + _hierarchical_bias_k(max_level) + ].shape[0] + num_classes = model.head.bias.shape[0] + if num_classes == finegrained_num_classes_pretrained == 10_000: + # Probably fine-tuning on iNat21 + logger.warn( + "Assuming that you pre-trained on iNat21 and are now fine-tuning on iNat21." + ) + state_dict["head.weight"] = state_dict[_hierarchical_weight_k(max_level)] + state_dict["head.bias"] = state_dict[_hierarchical_bias_k(max_level)] + else: + for i in range(max_level): + del state_dict[_hierarchical_weight_k(i)] + del state_dict[_hierarchical_bias_k(i)] + logger.warning( + "Error in loading classifier head, using default initialization." + ) + elif not pretrained_hierarchical and current_hierarchical: + # UNTESTED + assert "head.bias" in state_dict, "Should have a single pre-trained linear head" + assert not hasattr( + model.head, "bias" + ), "Should not have a single random linear head" + + # Delete the head.bias and head.weight keys then do nothing since the linear + # layer is already correctly initialized from scratch so it can fine-tune. + del state_dict["head.weight"] + del state_dict["head.bias"] + + elif pretrained_hierarchical and current_hierarchical: + assert ( + "head.bias" not in state_dict + ), "Should not have a single pre-trained linear head" + assert not hasattr( + model.head, "bias" + ), "Should not have a single random linear head" + + # Check if the two models have the exact same number of levels, and the same + # number of classes in each level + matches = True + level = 0 + while matches and _hierarchical_bias_k(level) in state_dict: + # Check that the current model has the right attribute + if level > len(model.head.heads): + matches = False + continue + + if not hasattr(model.head.heads[level], "bias"): + matches = False + continue + + if ( + model.head.heads[level].bias.shape + != state_dict[_hierarchical_bias_k(level)].shape + ): + matches = False + continue + + if ( + model.head.heads[level].weight.shape + != state_dict[_hierarchical_weight_k(level)].shape + ): + matches = False + continue + + level += 1 + + if not matches: + # UNTESTED + logger.warning( + "Not using pre-trained hierarchical head because the shapes do not match." + ) + # Delete the keys from the state dict because the pre-trained model and + # the current model do not match in size. + for i in range(level): + del state_dict[_hierarchical_weight_k(i)] + del state_dict[_hierarchical_bias_k(i)] + else: + logger.info("Using pre-trained hierarchical head.") + + def load_pretrained(config, model, logger): logger.info( f"==============> Loading weight {config.MODEL.PRETRAINED} for fine-tuning......" @@ -137,29 +287,15 @@ def load_pretrained(config, model, logger): ) state_dict[k] = absolute_pos_embed_pretrained_resized - # check classifier, if not match, then re-init classifier to zero - head_bias_pretrained = state_dict["head.bias"] - Nc1 = head_bias_pretrained.shape[0] - Nc2 = model.head.bias.shape[0] - if Nc1 != Nc2: - if Nc1 == 21841 and Nc2 == 1000: - logger.info("loading ImageNet-22K weight to ImageNet-1K ......") - map22kto1k_path = f"data/map22kto1k.txt" - with open(map22kto1k_path) as f: - map22kto1k = f.readlines() - map22kto1k = [int(id22k.strip()) for id22k in map22kto1k] - state_dict["head.weight"] = state_dict["head.weight"][map22kto1k, :] - state_dict["head.bias"] = state_dict["head.bias"][map22kto1k] - else: - torch.nn.init.constant_(model.head.bias, 0.0) - torch.nn.init.constant_(model.head.weight, 0.0) - del state_dict["head.weight"] - del state_dict["head.bias"] - logger.warning( - f"Error in loading classifier head, re-init classifier head to 0" - ) + handle_linear_head(config, model, state_dict, logger) msg = model.load_state_dict(state_dict, strict=False) + for key in msg.missing_keys: + assert ( + "relative_coords_table" in key + or "relative_position_index" in key + or "attn_mask" in key + ), f"Should only reinitialize relative positional information, not '{key}'" logger.warning(msg) logger.info(f"=> loaded successfully '{config.MODEL.PRETRAINED}'") From 11d11d92ed87e163a677b5401cca8132b1c6b63e Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 18:34:16 +0000 Subject: [PATCH 22/47] Add hierarchical fine-tuning config --- ...92to256_inat21_hierarchical_lr1.25_ft.yaml | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml diff --git a/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml b/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml new file mode 100644 index 00000000..906961e0 --- /dev/null +++ b/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml @@ -0,0 +1,34 @@ +DATA: + DATASET: inat21 + NUM_WORKERS: 32 + BATCH_SIZE: 16 + IMG_SIZE: 256 + DATA_PATH: /mnt/10tb/data/inat21/resize-256 +MODEL: + TYPE: swinv2 + NAME: swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft + PRETRAINED: /mnt/10tb/models/swinv2_base_patch4_window12_192_inat21_hierarchical_lr1.25_v0_epoch_89.pth + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 128 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 4, 8, 16, 32 ] + WINDOW_SIZE: 16 + PRETRAINED_WINDOW_SIZES: [ 12, 12, 12, 6 ] +TRAIN: + # Global batch size of 1024 + ACCUMULATION_STEPS: 8 + EPOCHS: 30 + WARMUP_EPOCHS: 5 + + # Should weight decay be this low? + # I don't think so, but I am sticking with the default for now. + WEIGHT_DECAY: 1.0e-8 + + BASE_LR: 2.0e-05 + WARMUP_LR: 2.0e-08 + MIN_LR: 2.0e-07 + + HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] + +HIERARCHICAL: true From 7c94a92c76bb3e5d72f6581ef0750b52ca27dfab Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 18:52:24 +0000 Subject: [PATCH 23/47] Adjust wandb panels --- main.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/main.py b/main.py index e562778d..a9f53402 100644 --- a/main.py +++ b/main.py @@ -359,15 +359,15 @@ def train_one_epoch( f"mem {memory_used:.0f}MB" ) stats = { - "train_batch_time": batch_time.val, - "train_loss": loss_meter.val, - "train_grad_norm": norm_meter.val, - "train_loss_scale": scaler_meter.val, + "train/batch_time": batch_time.val, + "train/loss": loss_meter.val, + "train/grad_norm": norm_meter.val, + "train/loss_scale": scaler_meter.val, "memory_mb": memory_used, - "learning_rate": config.TRAIN.BASE_LR, + "train/learning_rate": config.TRAIN.BASE_LR, } if lr_scheduler is not None: - stats["learning_rate"] = lr_scheduler.get_update_values( + stats["train/learning_rate"] = lr_scheduler.get_update_values( # Copied from line 326 (epoch * num_steps + idx) // config.TRAIN.ACCUMULATION_STEPS @@ -381,7 +381,7 @@ def train_one_epoch( logger.info( f"EPOCH {epoch} training took {datetime.timedelta(seconds=int(epoch_time))}" ) - wandb_writer.log({"train_time": epoch_time, "epoch": epoch}, epoch) + wandb_writer.log({"train/epoch_time": epoch_time, "epoch": epoch}, epoch) @torch.no_grad() @@ -435,9 +435,9 @@ def validate(config, data_loader, model, epoch): logger.info(f" * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}") wandb_writer.log( { - "val_acc1": acc1_meter.avg, - "val_acc5": acc5_meter.avg, - "val_loss": loss_meter.avg, + "val/acc1": acc1_meter.avg, + "val/acc5": acc5_meter.avg, + "val/loss": loss_meter.avg, "epoch": epoch, }, ) From b0f6dd7c06aac347ed6f5e9552975d236713daef Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 18:52:37 +0000 Subject: [PATCH 24/47] Include .python-version --- .gitignore | 3 --- 1 file changed, 3 deletions(-) diff --git a/.gitignore b/.gitignore index 8e3e1076..71b03bbb 100644 --- a/.gitignore +++ b/.gitignore @@ -87,9 +87,6 @@ target/ profile_default/ ipython_config.py -# pyenv -.python-version - # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. # However, in case of collaboration, if having platform-specific dependencies or dependencies From 7daba5899e06df87eea71573e8d06d55b1f030ea Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 19:08:52 +0000 Subject: [PATCH 25/47] Improve wandb logging --- logger.py | 23 ++++++++++++++++++++--- main.py | 2 +- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/logger.py b/logger.py index a557d2d3..d6ae3880 100644 --- a/logger.py +++ b/logger.py @@ -93,10 +93,27 @@ def init(self, config): wandb.init( config=config, project="hierarchical-vision", dir="./runs", resume=True ) + # Validation metrics + wandb.define_metric( + "val/loss", step_metric="epoch", summary="best", objective="max" + ) + wandb.define_metric( + "val/acc1", step_metric="epoch", summary="best", objective="max" + ) + wandb.define_metric( + "val/acc5", step_metric="epoch", summary="best", objective="max" + ) + + # Training metrics + wandb.define_metric("train/batch_time", step_metric="step", summary="last") + wandb.define_metric("train/grad_norm", step_metric="step", summary="last") + wandb.define_metric("train/loss", step_metric="step", summary="last") + wandb.define_metric("train/loss_scale", step_metric="step", summary="last") + wandb.define_metric("train/learning_rate", step_metric="step", summary="last") + wandb.define_metric("train/epoch_time", step_metric="epoch", summary="last") - wandb.define_metric("val/loss", step_metric="epoch", summary="min") - wandb.define_metric("val/acc1", step_metric="epoch", summary="max") - wandb.define_metric("val/acc5", step_metric="epoch", summary="max") + # Other metrics + wandb.define_metric("memory_mb", summary="max") def log(self, dct): if self.rank != 0: diff --git a/main.py b/main.py index a9f53402..80fc8110 100644 --- a/main.py +++ b/main.py @@ -360,8 +360,8 @@ def train_one_epoch( ) stats = { "train/batch_time": batch_time.val, - "train/loss": loss_meter.val, "train/grad_norm": norm_meter.val, + "train/loss": loss_meter.val, "train/loss_scale": scaler_meter.val, "memory_mb": memory_used, "train/learning_rate": config.TRAIN.BASE_LR, From 1ccb788dda11f9a0988851539e202d2c70cc41ff Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 19:09:07 +0000 Subject: [PATCH 26/47] Add AWS helpers --- README.md | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index a2dc800c..7fc51e1a 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,7 @@ pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1+cu1 Python packages ```sh -pip install matplotlib yacs timm einops black isort flake8 flake8-bugbear termcolor tensorboard preface opencv-python +pip install matplotlib yacs timm einops black isort flake8 flake8-bugbear termcolor wandb preface opencv-python ``` 2. Install Apex @@ -71,3 +71,25 @@ Use your root data folder and your size of choice. python -m data.inat preprocess /mnt/10tb/data/inat21/ val resize 224 python -m data.inat preprocess /mnt/10tb/data/inat21/ train resize 224 ``` + +5. Login to Wandb + +``` +wandb login +``` + +## AWS Helpers + +Uninstall v1 of awscli: + +``` +sudo /usr/local/bin/pip uninstall awscli +``` + +Install v2: +``` +cd ~/pkg +curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" +unzip awscliv2.zip +./aws/install --bin-dir ~/.local/bin --install-dir ~/.local/aws-cli +``` From e865fc82b2b4e24e670c236a793689a2f5b968e9 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 19:22:17 +0000 Subject: [PATCH 27/47] Update wandb logging --- logger.py | 4 +++- main.py | 7 +++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/logger.py b/logger.py index d6ae3880..9185a1eb 100644 --- a/logger.py +++ b/logger.py @@ -107,10 +107,12 @@ def init(self, config): # Training metrics wandb.define_metric("train/batch_time", step_metric="step", summary="last") wandb.define_metric("train/grad_norm", step_metric="step", summary="last") - wandb.define_metric("train/loss", step_metric="step", summary="last") + wandb.define_metric("train/batch_loss", step_metric="step", summary="last") wandb.define_metric("train/loss_scale", step_metric="step", summary="last") wandb.define_metric("train/learning_rate", step_metric="step", summary="last") + wandb.define_metric("train/epoch_time", step_metric="epoch", summary="last") + wandb.define_metric("train/loss", step_metric="epoch", summary="last") # Other metrics wandb.define_metric("memory_mb", summary="max") diff --git a/main.py b/main.py index 80fc8110..55731a4c 100644 --- a/main.py +++ b/main.py @@ -360,8 +360,8 @@ def train_one_epoch( ) stats = { "train/batch_time": batch_time.val, + "train/batch_loss": loss_meter.val, "train/grad_norm": norm_meter.val, - "train/loss": loss_meter.val, "train/loss_scale": scaler_meter.val, "memory_mb": memory_used, "train/learning_rate": config.TRAIN.BASE_LR, @@ -381,7 +381,10 @@ def train_one_epoch( logger.info( f"EPOCH {epoch} training took {datetime.timedelta(seconds=int(epoch_time))}" ) - wandb_writer.log({"train/epoch_time": epoch_time, "epoch": epoch}, epoch) + wandb_writer.log( + {"train/epoch_time": epoch_time, "train/loss": loss_meter.avg, "epoch": epoch}, + epoch, + ) @torch.no_grad() From b6db1b2308d3733b5cff8f1b6b0a412a5a73cc03 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 19:15:12 +0000 Subject: [PATCH 28/47] Fix wandb auto-resuming --- logger.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/logger.py b/logger.py index 9185a1eb..880cc641 100644 --- a/logger.py +++ b/logger.py @@ -91,7 +91,10 @@ def init(self, config): return wandb.init( - config=config, project="hierarchical-vision", dir="./runs", resume=True + config=config, + project="hierarchical-vision", + dir="./runs", + resume=bool(config.MODEL.RESUME), ) # Validation metrics wandb.define_metric( From a356abf55e9196e822f4c776e45b91005ae1b889 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 19:28:28 +0000 Subject: [PATCH 29/47] Clean configs --- configs/swinv2/a6000.yaml | 131 ------------------ configs/swinv2/large_inat21_a6000.yaml | 31 ----- configs/swinv2/large_inat21_constant_lr.yaml | 30 ---- ...e_patch4_window12_inat21_hierarchical.yaml | 26 ---- ...window12_inat21_hierarchical_debugging.yml | 40 ------ ...indow12_192_inat21_hierarchical_lr2.5.yaml | 2 +- ...192to256_inat21_hierarchical_lr2.5_ft.yaml | 34 +++++ 7 files changed, 35 insertions(+), 259 deletions(-) delete mode 100644 configs/swinv2/a6000.yaml delete mode 100644 configs/swinv2/large_inat21_a6000.yaml delete mode 100644 configs/swinv2/large_inat21_constant_lr.yaml delete mode 100644 configs/swinv2/large_patch4_window12_inat21_hierarchical.yaml delete mode 100644 configs/swinv2/large_patch4_window12_inat21_hierarchical_debugging.yml create mode 100644 configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft.yaml diff --git a/configs/swinv2/a6000.yaml b/configs/swinv2/a6000.yaml deleted file mode 100644 index d9930c4a..00000000 --- a/configs/swinv2/a6000.yaml +++ /dev/null @@ -1,131 +0,0 @@ -AMP_ENABLE: true -AMP_OPT_LEVEL: '' -AUG: - AUTO_AUGMENT: rand-m9-mstd0.5-inc1 - COLOR_JITTER: 0.4 - CUTMIX: 1.0 - CUTMIX_MINMAX: null - MIXUP: 0.8 - MIXUP_MODE: batch - MIXUP_PROB: 1.0 - MIXUP_SWITCH_PROB: 0.5 - RECOUNT: 1 - REMODE: pixel - REPROB: 0.25 -BASE: -- '' -DATA: - BATCH_SIZE: 128 - CACHE_MODE: part - DATASET: inat21 - DATA_PATH: /research/nfs_su_809/cv_datasets/inat21/train_val_224 - IMG_SIZE: 224 - INTERPOLATION: bicubic - NUM_WORKERS: 8 - PIN_MEMORY: true - ZIP_MODE: false -EVAL_MODE: false -LOCAL_RANK: 0 -MODEL: - DROP_PATH_RATE: 0.2 - DROP_RATE: 0.0 - LABEL_SMOOTHING: 0.1 - NAME: swinv2_large_patch4_window7_224_inat21 - NUM_CLASSES: 1000 - PRETRAINED: '' - RESUME: '' - SWIN: - APE: false - DEPTHS: - - 2 - - 2 - - 6 - - 2 - EMBED_DIM: 96 - IN_CHANS: 3 - MLP_RATIO: 4.0 - NUM_HEADS: - - 3 - - 6 - - 12 - - 24 - PATCH_NORM: true - PATCH_SIZE: 4 - QKV_BIAS: true - QK_SCALE: null - WINDOW_SIZE: 7 - SWINV2: - APE: false - DEPTHS: - - 2 - - 2 - - 18 - - 2 - EMBED_DIM: 192 - IN_CHANS: 3 - MLP_RATIO: 4.0 - NUM_HEADS: - - 6 - - 12 - - 24 - - 48 - PATCH_NORM: true - PATCH_SIZE: 4 - PRETRAINED_WINDOW_SIZES: - - 0 - - 0 - - 0 - - 0 - QKV_BIAS: true - WINDOW_SIZE: 7 - SWIN_MLP: - APE: false - DEPTHS: - - 2 - - 2 - - 6 - - 2 - EMBED_DIM: 96 - IN_CHANS: 3 - MLP_RATIO: 4.0 - NUM_HEADS: - - 3 - - 6 - - 12 - - 24 - PATCH_NORM: true - PATCH_SIZE: 4 - WINDOW_SIZE: 7 - TYPE: swinv2 -OUTPUT: /home/stevens.994/projects/Swin-Transformer/model-weights/swinv2_large_patch4_window7_224_inat21/test2 -PRINT_FREQ: 10 -SAVE_FREQ: 1 -SEED: 0 -TAG: test2 -TEST: - CROP: true - SEQUENTIAL: false -THROUGHPUT_MODE: false -TRAIN: - ACCUMULATION_STEPS: 1 - AUTO_RESUME: true - BASE_LR: 0.000125 - CLIP_GRAD: 5.0 - EPOCHS: 180 - LR_SCHEDULER: - DECAY_EPOCHS: 30 - DECAY_RATE: 0.1 - NAME: cosine - MIN_LR: 1.25e-06 - OPTIMIZER: - BETAS: - - 0.9 - - 0.999 - EPS: 1.0e-08 - MOMENTUM: 0.9 - NAME: adamw - START_EPOCH: 0 - USE_CHECKPOINT: false - WARMUP_EPOCHS: 5 - WARMUP_LR: 1.25e-07 - WEIGHT_DECAY: 0.1 diff --git a/configs/swinv2/large_inat21_a6000.yaml b/configs/swinv2/large_inat21_a6000.yaml deleted file mode 100644 index cf448a03..00000000 --- a/configs/swinv2/large_inat21_a6000.yaml +++ /dev/null @@ -1,31 +0,0 @@ -DATA: - DATASET: inat21 - IMG_SIZE: 224 - BATCH_SIZE: 32 - DATA_PATH: /mnt/10tb/data/inat21/resize-224 - NUM_WORKERS: 32 -MODEL: - TYPE: swinv2 - NAME: swinv2_large_resize_224_inat21 - DROP_PATH_RATE: 0.2 - SWINV2: - EMBED_DIM: 192 - DEPTHS: [ 2, 2, 18, 2 ] - NUM_HEADS: [ 6, 12, 24, 48 ] - WINDOW_SIZE: 7 -TRAIN: - EPOCHS: 500 - WARMUP_EPOCHS: 5 - WEIGHT_DECAY: 0.1 - BASE_LR: 1.25e-4 - WARMUP_LR: 1.25e-7 - CLIP_GRAD: 5.0 - # N_GPU * BATCH_SIZE * ACCUMULATION_STEPS = 512 - ACCUMULATION_STEPS: 2 - - LR_SCHEDULER: - NAME: constant - -# Only save checkpoint every 5 epochs -SAVE_FREQ: 5 - diff --git a/configs/swinv2/large_inat21_constant_lr.yaml b/configs/swinv2/large_inat21_constant_lr.yaml deleted file mode 100644 index bfce8b04..00000000 --- a/configs/swinv2/large_inat21_constant_lr.yaml +++ /dev/null @@ -1,30 +0,0 @@ -DATA: - DATASET: inat21 - IMG_SIZE: 192 - BATCH_SIZE: 32 - DATA_PATH: /mnt/10tb/data/inat21/resize-192 - NUM_WORKERS: 32 -MODEL: - TYPE: swinv2 - NAME: swinv2_large_resize_192_inat21 - DROP_PATH_RATE: 0.2 - SWINV2: - EMBED_DIM: 192 - DEPTHS: [ 2, 2, 18, 2 ] - NUM_HEADS: [ 6, 12, 24, 48 ] - WINDOW_SIZE: 12 -TRAIN: - EPOCHS: 240 - WARMUP_EPOCHS: 20 - WEIGHT_DECAY: 0.1 - BASE_LR: 1.0e-3 - WARMUP_LR: 1.0e-7 - CLIP_GRAD: 5.0 - # N_GPU * BATCH_SIZE * ACCUMULATION_STEPS ~= 4096 - ACCUMULATION_STEPS: 16 - - LR_SCHEDULER: - NAME: cosine - -# Only save checkpoint every 5 epochs -SAVE_FREQ: 5 diff --git a/configs/swinv2/large_patch4_window12_inat21_hierarchical.yaml b/configs/swinv2/large_patch4_window12_inat21_hierarchical.yaml deleted file mode 100644 index 3d6761bf..00000000 --- a/configs/swinv2/large_patch4_window12_inat21_hierarchical.yaml +++ /dev/null @@ -1,26 +0,0 @@ -DATA: - DATASET: inat21 - IMG_SIZE: 192 - BATCH_SIZE: 32 - DATA_PATH: /mnt/10tb/data/inat21/resize-192 - NUM_WORKERS: 32 -MODEL: - TYPE: swinv2 - NAME: swinv2_large_patch4_window12_resize_192_inat21 - DROP_PATH_RATE: 0.2 - SWINV2: - EMBED_DIM: 192 - DEPTHS: [ 2, 2, 18, 2 ] - NUM_HEADS: [ 6, 12, 24, 48 ] - WINDOW_SIZE: 12 -TRAIN: - EPOCHS: 90 - WARMUP_EPOCHS: 5 - WEIGHT_DECAY: 0.1 - BASE_LR: 1.25e-4 # 4096 batch-size - WARMUP_LR: 1.25e-7 - MIN_LR: 1.25e-6 - HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] - # We want N_GPU * BATCH_SIZE * ACCUMULATION_STEPS ~= 1024 - ACCUMULATION_STEPS: 4 -HIERARHICAL: true diff --git a/configs/swinv2/large_patch4_window12_inat21_hierarchical_debugging.yml b/configs/swinv2/large_patch4_window12_inat21_hierarchical_debugging.yml deleted file mode 100644 index 9462ab9d..00000000 --- a/configs/swinv2/large_patch4_window12_inat21_hierarchical_debugging.yml +++ /dev/null @@ -1,40 +0,0 @@ -DATA: - DATASET: inat21 - IMG_SIZE: 192 - BATCH_SIZE: 32 - DATA_PATH: /mnt/10tb/data/inat21/resize-192 - NUM_WORKERS: 32 -AUG: - COLOR_JITTER: 0.0 - AUTO_AUGMENT: "none" - REPROB: 0.0 - RECOUNT: 0 - MIXUP: 0.0 - CUTMIX: 0.0 - MIXUP_PROB: 0.0 -MODEL: - TYPE: swinv2 - NAME: swinv2_large_patch4_window12_resize_192_inat21_debugging - SWINV2: - EMBED_DIM: 192 - DEPTHS: [ 2, 2, 18, 2 ] - NUM_HEADS: [ 6, 12, 24, 48 ] - WINDOW_SIZE: 12 - # Regularization - DROP_RATE: 0.0 - DROP_PATH_RATE: 0.0 - LABEL_SMOOTHING: 0.0 -TRAIN: - OVERFIT_BATCHES: 8 - EPOCHS: 90 - WARMUP_EPOCHS: 0 - WEIGHT_DECAY: 0.0 - BASE_LR: 1.25e-4 # 4096 batch-size - WARMUP_LR: 1.25e-4 - MIN_LR: 1.25e-4 - HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] - # We want N_GPU * BATCH_SIZE * ACCUMULATION_STEPS ~= 1024 - ACCUMULATION_STEPS: 4 - LR_SCHEDULER: - NAME: "none" -HIERARHICAL: true diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml index 9deae81e..88456b7d 100644 --- a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml +++ b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml @@ -30,4 +30,4 @@ TRAIN: MIN_LR: 2.5e-6 HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] -HIERARHICAL: true +HIERARCHICAL: true diff --git a/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft.yaml b/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft.yaml new file mode 100644 index 00000000..ee5735da --- /dev/null +++ b/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft.yaml @@ -0,0 +1,34 @@ +DATA: + DATASET: inat21 + NUM_WORKERS: 32 + BATCH_SIZE: 16 + IMG_SIZE: 256 + DATA_PATH: /mnt/10tb/data/inat21/resize-256 +MODEL: + TYPE: swinv2 + NAME: swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft + PRETRAINED: /mnt/10tb/models/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5_v0_epoch_89.pth + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 128 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 4, 8, 16, 32 ] + WINDOW_SIZE: 16 + PRETRAINED_WINDOW_SIZES: [ 12, 12, 12, 6 ] +TRAIN: + # Global batch size of 1024 + ACCUMULATION_STEPS: 8 + EPOCHS: 30 + WARMUP_EPOCHS: 5 + + # Should weight decay be this low? + # I don't think so, but I am sticking with the default for now. + WEIGHT_DECAY: 1.0e-8 + + BASE_LR: 2.0e-05 + WARMUP_LR: 2.0e-08 + MIN_LR: 2.0e-07 + + HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] + +HIERARCHICAL: true From 181178b23c24e51fcc81082d72e88381142e98af Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Sun, 30 Oct 2022 21:01:04 +0000 Subject: [PATCH 30/47] Fix logging error --- main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/main.py b/main.py index 55731a4c..b6311cf2 100644 --- a/main.py +++ b/main.py @@ -383,7 +383,6 @@ def train_one_epoch( ) wandb_writer.log( {"train/epoch_time": epoch_time, "train/loss": loss_meter.avg, "epoch": epoch}, - epoch, ) From 0c418901603d29304a94c08c41f51b2b2250f30d Mon Sep 17 00:00:00 2001 From: Samuel Stevens Date: Mon, 31 Oct 2022 14:36:29 -0400 Subject: [PATCH 31/47] Add default config --- ...winv2_base_patch4_window12_192_inat21.yaml | 27 +++++++++++++++++++ scripts/train.fish | 14 +++++++--- 2 files changed, 38 insertions(+), 3 deletions(-) create mode 100644 configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml b/configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml new file mode 100644 index 00000000..d283f68a --- /dev/null +++ b/configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml @@ -0,0 +1,27 @@ +DATA: + DATASET: inat21 + IMG_SIZE: 192 + DATA_PATH: /research/nfs_su_809/cv_datasets/inat21/train_val_192 + NUM_WORKERS: 32 + BATCH_SIZE: 256 +MODEL: + TYPE: swinv2 + NAME: swinv2_base_patch4_window12_192_inat21 + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 128 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 4, 8, 16, 32 ] + WINDOW_SIZE: 12 +TRAIN: + # Want a global batch size of 2048 because SwinV2 was trained on 16 V100s with batch size 128 (I think) + # But we are going to use a global batch size of 1024 because it's faster (throughput). + ACCUMULATION_STEPS: 1 + + # We are using limited epochs based on pre-training configs for imagenet22k + # Then we will pre-train on 256x256 for 30 epochs + EPOCHS: 90 + WARMUP_EPOCHS: 5 + WEIGHT_DECAY: 0.1 + +SAVE_FREQ: 4 diff --git a/scripts/train.fish b/scripts/train.fish index b7fcb334..a0c1c351 100644 --- a/scripts/train.fish +++ b/scripts/train.fish @@ -8,6 +8,7 @@ function usage echo " -h/--help print this message" echo " --config which config file to use (should be YAML)" echo " --debug run with only one process rather than 8 so pdb is useable." + echo " --nprocs number of processes" echo " --tag tag for the run (I use v0, v1, etc)" echo " --venv path to the virtual environment (default ./venv/)" end @@ -15,7 +16,8 @@ end set -l options (fish_opt --short h --long help) set -a options (fish_opt --short c --long config --required-val --long-only) set -a options (fish_opt --short t --long tag --required-val --long-only) -set -a options (fish_opt --short v --long venv --long-only) +set -a options (fish_opt --short p --long nprocs --required-val --long-only) +set -a options (fish_opt --short v --long venv --required-val --long-only) set -a options (fish_opt --short d --long debug --long-only) argparse $options -- $argv @@ -40,14 +42,20 @@ if not test -z $_flag_venv set venv $_flag_venv end +set nprocs 8 +if not test -z $_flag_nprocs + set nprocs $_flag_nprocs +end + set launcher $venv/bin/torchrun -set launcher_args --nproc_per_node 8 --master_port 12345 +set launcher_args --nproc_per_node $nprocs if not test -z $_flag_debug - set launcher_args --nproc_per_node 1 --master_port 12345 + set launcher_args --nproc_per_node 1 end $launcher $launcher_args \ + --master_port 12345 \ main.py \ --cfg $_flag_config \ --output runs \ From ee1eecc2fb6cce4e4d83d6da07f96b3ede2a256c Mon Sep 17 00:00:00 2001 From: Samuel Stevens Date: Mon, 31 Oct 2022 14:38:39 -0400 Subject: [PATCH 32/47] Fix std/mean for the last time! --- data/constants.py | 4 ++++ data/inat/__main__.py | 6 +++--- data/inat/datasets.py | 2 +- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/data/constants.py b/data/constants.py index 5bb8b2eb..001823ae 100644 --- a/data/constants.py +++ b/data/constants.py @@ -6,6 +6,10 @@ torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), ), + "/research/nfs_su_809/cv_datasets/inat21/train_val_192": ( + torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), + torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), + ), "/mnt/10tb/data/inat21/resize-224": ( torch.tensor([0.23762744665145874, 0.2292044311761856, 0.24757201969623566]), torch.tensor([0.4632636606693268, 0.48004215955734253, 0.37622377276420593]), diff --git a/data/inat/__main__.py b/data/inat/__main__.py index 0feeb2df..c1c0883c 100644 --- a/data/inat/__main__.py +++ b/data/inat/__main__.py @@ -9,13 +9,13 @@ def preprocess_cli(args): def normalize_cli(args): - std, mean = datasets.load_statistics(args.directory) + mean, std = datasets.load_statistics(args.directory) print("Add this to a constants.py file:") print( f""" "{args.directory}": ( - torch.tensor({std.tolist()}), torch.tensor({mean.tolist()}), + torch.tensor({std.tolist()}), ),""" ) @@ -50,7 +50,7 @@ def num_classes_cli(args): # Normalize normalize_parser = subparsers.add_parser( - "normalize", help="Measure std. dev. and mean of dataset." + "normalize", help="Measure mean and std of dataset." ) normalize_parser.add_argument("directory", help="Data folder") normalize_parser.set_defaults(func=normalize_cli) diff --git a/data/inat/datasets.py b/data/inat/datasets.py index dcf2e356..b0676896 100644 --- a/data/inat/datasets.py +++ b/data/inat/datasets.py @@ -37,7 +37,7 @@ def load_statistics(directory): var = total_squared / divisor - torch.mul(mean, mean) std = torch.sqrt(var) - return std, mean + return mean, std class Inat21: From f6a748b635b0992372d59f64fcfe52b794e0a44c Mon Sep 17 00:00:00 2001 From: Samuel Stevens Date: Mon, 31 Oct 2022 14:38:57 -0400 Subject: [PATCH 33/47] Default hierarchical multitask config --- ..._window12_192_inat21_hierarchical_lr5.yaml | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr5.yaml diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr5.yaml b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr5.yaml new file mode 100644 index 00000000..75a9bc99 --- /dev/null +++ b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr5.yaml @@ -0,0 +1,29 @@ +DATA: + DATASET: inat21 + IMG_SIZE: 192 + DATA_PATH: /research/nfs_su_809/cv_datasets/inat21/train_val_192 + NUM_WORKERS: 32 + BATCH_SIZE: 128 +MODEL: + TYPE: swinv2 + NAME: swinv2_base_patch4_window12_192_inat21_hierarchical_lr5 + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 128 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 4, 8, 16, 32 ] + WINDOW_SIZE: 12 +TRAIN: + # Want a global batch size of 2048 because SwinV2 was trained on 16 V100s with batch size 128 (I think) + # But we are going to use a global batch size of 1024 because it's faster (throughput). + ACCUMULATION_STEPS: 2 + + # We are using limited epochs based on pre-training configs for imagenet22k + # Then we will pre-train on 256x256 for 30 epochs + EPOCHS: 90 + WARMUP_EPOCHS: 5 + WEIGHT_DECAY: 0.1 + HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] + +SAVE_FREQ: 4 +HIERARCHICAL: true From 9e7b4c15da467c17ae1b9808578c8d55eeaca76c Mon Sep 17 00:00:00 2001 From: Samuel Stevens Date: Mon, 31 Oct 2022 15:16:56 -0400 Subject: [PATCH 34/47] updated experiment name to funky-banana-192 --- ...winv2_base_patch4_window12_192_inat21.yaml | 2 +- docs/experiments/funky-banana.md | 19 +++++++++++++++++++ docs/experiments/groovy-grape.md | 18 ++++++++++++++++++ docs/experiments/index.md | 5 +++++ 4 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 docs/experiments/funky-banana.md create mode 100644 docs/experiments/groovy-grape.md create mode 100644 docs/experiments/index.md diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml b/configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml index d283f68a..63f05fac 100644 --- a/configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml +++ b/configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml @@ -6,7 +6,7 @@ DATA: BATCH_SIZE: 256 MODEL: TYPE: swinv2 - NAME: swinv2_base_patch4_window12_192_inat21 + NAME: funky-banana-192 DROP_PATH_RATE: 0.2 SWINV2: EMBED_DIM: 128 diff --git a/docs/experiments/funky-banana.md b/docs/experiments/funky-banana.md new file mode 100644 index 00000000..5e00f16d --- /dev/null +++ b/docs/experiments/funky-banana.md @@ -0,0 +1,19 @@ +# Funky Banana + +This experiment is the default swin-v2-base applied to iNat 21 192x192. +We train for 90 epochs at 192x192, then tune for 30 epochs at 256x256. + +```yaml +configs: +- configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml +- configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21.yaml +codename: funky-banana +``` + +## Log + +I initialized training on strawberry0 on 4x A6000 servers. + +I decided to use the A6000 servers for 256x256 tuning, so I am moving the latest checkpoint to S3, then cloning it back to an 8x V100 server to finish training. +I am storing the 40th checkpoint on S3 as funky-banana-192-epoch40.pth. + diff --git a/docs/experiments/groovy-grape.md b/docs/experiments/groovy-grape.md new file mode 100644 index 00000000..6b219363 --- /dev/null +++ b/docs/experiments/groovy-grape.md @@ -0,0 +1,18 @@ +# Groovy Grape + +This experiment trains swin-v2-base from scratch on iNat21 using the multitask objective using 1/4 of the default learning rate, which works out to 1.25e-4. +It also does 90 epochs at 192, then 30 epochs at 256. + +```yaml +configs: +- configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr1.25.yaml +- configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml +codename: groovy-grape +``` + +## Log + +This model trained the first 90 on 8x V100, and did 8/30 epochs at 256 on 8x V100. +I am storing the 8th checkpoint on S3 as groovy-grape-256-epoch8.pth. +Now it is stored at /local/scratch/stevens.994/hierarchical-vision/groovy-grape-256/v0 +It was originally haunted-broomstick on wandb, but is now groovy-grape-256. diff --git a/docs/experiments/index.md b/docs/experiments/index.md new file mode 100644 index 00000000..1441fbab --- /dev/null +++ b/docs/experiments/index.md @@ -0,0 +1,5 @@ +# Experiment Index + +This file has a lightweight index of all the experiments I'm running so you can go to the file if necessary. + +[funky-banana](funky-banana.md): Default swin-v2-base on iNat21 From 9cdafe63aeb6539a36f8f510234ea9f33fe8295f Mon Sep 17 00:00:00 2001 From: Samuel Stevens Date: Mon, 31 Oct 2022 15:21:05 -0400 Subject: [PATCH 35/47] Update experiment name tracking --- ...ch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml | 2 +- logger.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml b/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml index 906961e0..2e88e62b 100644 --- a/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml +++ b/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml @@ -6,8 +6,8 @@ DATA: DATA_PATH: /mnt/10tb/data/inat21/resize-256 MODEL: TYPE: swinv2 - NAME: swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft PRETRAINED: /mnt/10tb/models/swinv2_base_patch4_window12_192_inat21_hierarchical_lr1.25_v0_epoch_89.pth + NAME: groovy-grape-256 DROP_PATH_RATE: 0.2 SWINV2: EMBED_DIM: 128 diff --git a/logger.py b/logger.py index 880cc641..ff3bf7c9 100644 --- a/logger.py +++ b/logger.py @@ -95,6 +95,7 @@ def init(self, config): project="hierarchical-vision", dir="./runs", resume=bool(config.MODEL.RESUME), + name=config.MODEL.NAME, ) # Validation metrics wandb.define_metric( From 2fd00e40baef1674a90b15e6190b25d8775d8f2b Mon Sep 17 00:00:00 2001 From: Samuel Stevens Date: Mon, 31 Oct 2022 18:07:22 -0400 Subject: [PATCH 36/47] Added additional fruity experiments --- ...indow12_192_inat21_hierarchical_lr2.5.yaml | 2 +- ..._window12_192_inat21_hierarchical_lr5.yaml | 2 +- ...192to256_inat21_hierarchical_lr2.5_ft.yaml | 2 +- docs/experiments/funky-banana.md | 2 +- docs/experiments/index.md | 3 +++ docs/experiments/outrageous-orange.md | 18 ++++++++++++++++++ docs/experiments/sweet-strawberry.md | 19 +++++++++++++++++++ 7 files changed, 44 insertions(+), 4 deletions(-) create mode 100644 docs/experiments/outrageous-orange.md create mode 100644 docs/experiments/sweet-strawberry.md diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml index 88456b7d..5bc4336f 100644 --- a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml +++ b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml @@ -6,7 +6,7 @@ DATA: BATCH_SIZE: 64 MODEL: TYPE: swinv2 - NAME: swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5 + NAME: sweet-strawberry-192 DROP_PATH_RATE: 0.2 SWINV2: EMBED_DIM: 128 diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr5.yaml b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr5.yaml index 75a9bc99..2a2d422f 100644 --- a/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr5.yaml +++ b/configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr5.yaml @@ -6,7 +6,7 @@ DATA: BATCH_SIZE: 128 MODEL: TYPE: swinv2 - NAME: swinv2_base_patch4_window12_192_inat21_hierarchical_lr5 + NAME: outrageous-orange-192 DROP_PATH_RATE: 0.2 SWINV2: EMBED_DIM: 128 diff --git a/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft.yaml b/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft.yaml index ee5735da..c3f6b5df 100644 --- a/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft.yaml +++ b/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft.yaml @@ -6,8 +6,8 @@ DATA: DATA_PATH: /mnt/10tb/data/inat21/resize-256 MODEL: TYPE: swinv2 - NAME: swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft PRETRAINED: /mnt/10tb/models/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5_v0_epoch_89.pth + NAME: sweet-strawberry-256 DROP_PATH_RATE: 0.2 SWINV2: EMBED_DIM: 128 diff --git a/docs/experiments/funky-banana.md b/docs/experiments/funky-banana.md index 5e00f16d..d5b88fde 100644 --- a/docs/experiments/funky-banana.md +++ b/docs/experiments/funky-banana.md @@ -16,4 +16,4 @@ I initialized training on strawberry0 on 4x A6000 servers. I decided to use the A6000 servers for 256x256 tuning, so I am moving the latest checkpoint to S3, then cloning it back to an 8x V100 server to finish training. I am storing the 40th checkpoint on S3 as funky-banana-192-epoch40.pth. - +It is now running as funky-banana-192 on 8x V100. diff --git a/docs/experiments/index.md b/docs/experiments/index.md index 1441fbab..065f5d84 100644 --- a/docs/experiments/index.md +++ b/docs/experiments/index.md @@ -3,3 +3,6 @@ This file has a lightweight index of all the experiments I'm running so you can go to the file if necessary. [funky-banana](funky-banana.md): Default swin-v2-base on iNat21 +[outrageous-orange](outrageous-orange.md): Hierarchical swin-v2-base on iNat21 +[sweet-strawberry](sweet-strawberry.md): Hierarchical swin-v2-base on iNat21 with 1/2 LR +[groovy-grape](groovy-grape.md): Hierarchical swin-v2-base on iNat21 with 1/4 LR diff --git a/docs/experiments/outrageous-orange.md b/docs/experiments/outrageous-orange.md new file mode 100644 index 00000000..091a06af --- /dev/null +++ b/docs/experiments/outrageous-orange.md @@ -0,0 +1,18 @@ +# Outrageous Orange + +This experiment is the swin-v2-base with hierarchical multitask objective applied to iNat 21 192x192. +We train for 90 epochs at 192x192, then tune for 30 epochs at 256x256. + +```yaml +configs: +- configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr5.yaml +codename: outrageous-orange +``` + +## Log + +I initialized training on strawberry0 on 4x A6000 servers. + +I decided to use the A6000 servers for 256x256 tuning, so I am moving the latest checkpoint to S3, then cloning it back to an 8x V100 server to finish training. +I am storing the 36th checkpoint on S3 as `outrageous-orange-192-epoch40.pth`. +It is now running as `outrageous-orange-192` on 8x V100. diff --git a/docs/experiments/sweet-strawberry.md b/docs/experiments/sweet-strawberry.md new file mode 100644 index 00000000..2aac60a4 --- /dev/null +++ b/docs/experiments/sweet-strawberry.md @@ -0,0 +1,19 @@ +# Sweet Strawberry + +This experiment trains swin-v2-base from scratch on iNat21 using the multitask objective using 1/2 of the default learning rate, which works out to 2.5e-4. +It also does 90 epochs at 192, then 30 epochs at 256. + +```yaml +configs: +- configs/swinv2/swinv2_base_patch4_window12_192_inat21_hierarchical_lr2.5.yaml +- configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr2.5_ft.yaml +codename: sweet-strawberry +``` + +## Log + +This model trained the first 90 on 8x V100, and did 15/30 epochs at 256 on 8x V100. +I am storing the 15th checkpoint on S3 as sweet-strawberry-256-epoch15.pth. +Now it is stored at /local/scratch/stevens.994/hierarchical-vision/sweet-strawberry-256/v0 +It was originally unearthly-moon on wandb, but is now sweet-strawberry-256. +It is running on 4x A6000. From 27870665c8ad469b5c60c7c9be8fdbe9a5b9245e Mon Sep 17 00:00:00 2001 From: Samuel Stevens Date: Tue, 1 Nov 2022 09:09:20 -0400 Subject: [PATCH 37/47] update experiment log --- docs/experiments/index.md | 3 +++ docs/experiments/outrageous-orange.md | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/experiments/index.md b/docs/experiments/index.md index 065f5d84..e230e7f8 100644 --- a/docs/experiments/index.md +++ b/docs/experiments/index.md @@ -3,6 +3,9 @@ This file has a lightweight index of all the experiments I'm running so you can go to the file if necessary. [funky-banana](funky-banana.md): Default swin-v2-base on iNat21 + [outrageous-orange](outrageous-orange.md): Hierarchical swin-v2-base on iNat21 + [sweet-strawberry](sweet-strawberry.md): Hierarchical swin-v2-base on iNat21 with 1/2 LR + [groovy-grape](groovy-grape.md): Hierarchical swin-v2-base on iNat21 with 1/4 LR diff --git a/docs/experiments/outrageous-orange.md b/docs/experiments/outrageous-orange.md index 091a06af..f42b43c8 100644 --- a/docs/experiments/outrageous-orange.md +++ b/docs/experiments/outrageous-orange.md @@ -14,5 +14,5 @@ codename: outrageous-orange I initialized training on strawberry0 on 4x A6000 servers. I decided to use the A6000 servers for 256x256 tuning, so I am moving the latest checkpoint to S3, then cloning it back to an 8x V100 server to finish training. -I am storing the 36th checkpoint on S3 as `outrageous-orange-192-epoch40.pth`. +I am storing the 36th checkpoint on S3 as `outrageous-orange-192-epoch36.pth`. It is now running as `outrageous-orange-192` on 8x V100. From 84fa2d979e5a3dbe01b2f33fb7f0b6296d3250d9 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Tue, 1 Nov 2022 15:38:08 +0000 Subject: [PATCH 38/47] add rough parsing script --- scripts/rough_parse.bash | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 scripts/rough_parse.bash diff --git a/scripts/rough_parse.bash b/scripts/rough_parse.bash new file mode 100644 index 00000000..d67fa3e7 --- /dev/null +++ b/scripts/rough_parse.bash @@ -0,0 +1,6 @@ +#!/usr/env bash + +cat $1 | rg '190/196.*Acc@1' | nl | rg ' 1\t' +cat $1 | rg '190/196.*Acc@1' | nl | rg ' 30\t' +cat $1 | rg '190/196.*Acc@1' | nl | rg ' 60\t' +cat $1 | rg '190/196.*Acc@1' | nl | rg ' 90\t' From e956b1f10b4c898ee719d3b69423f902ae13acab Mon Sep 17 00:00:00 2001 From: Samuel Stevens Date: Tue, 1 Nov 2022 23:22:16 -0400 Subject: [PATCH 39/47] Rename config options --- .gitignore | 5 ++ README.md | 8 +- config.py | 58 ++++++++++--- .../funky-banana-192.yaml} | 10 ++- .../groovy-grape-256.yaml} | 14 ++-- data/build.py | 10 +-- data/constants.py | 6 +- docs/experiments/funky-banana.md | 2 +- logger.py | 30 +++---- main.py | 23 ++++-- main_moe.py | 4 +- main_simmim_ft.py | 4 +- main_simmim_pt.py | 4 +- scripts/generate_wandb_id.py | 3 + scripts/train.fish | 82 ++++++++++++++----- utils.py | 2 +- 16 files changed, 188 insertions(+), 77 deletions(-) rename configs/{swinv2/swinv2_base_patch4_window12_192_inat21.yaml => hierarchical-vision-project/funky-banana-192.yaml} (85%) rename configs/{swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml => hierarchical-vision-project/groovy-grape-256.yaml} (72%) create mode 100644 scripts/generate_wandb_id.py diff --git a/.gitignore b/.gitignore index 71b03bbb..9317223a 100644 --- a/.gitignore +++ b/.gitignore @@ -134,4 +134,9 @@ dmypy.json # Apex apex/ +# Logging runs/ +wandb/ + +# host-specific values +scripts/env.fish diff --git a/README.md b/README.md index 7fc51e1a..065e4bb9 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ python setup.py install We use the iNat21 dataseta available on [GitHub](https://github.com/visipedia/inat_comp/tree/master/2021) ``` +cd /mnt/10tb mkdir -p data/inat21 cd data/inat21 mkdir compressed raw @@ -68,8 +69,11 @@ mv train ../raw/ Use your root data folder and your size of choice. ``` -python -m data.inat preprocess /mnt/10tb/data/inat21/ val resize 224 -python -m data.inat preprocess /mnt/10tb/data/inat21/ train resize 224 +export DATA_DIR=/mnt/10tb/data/inat21/ +python -m data.inat preprocess $DATA_DIR val resize 192 +python -m data.inat preprocess $DATA_DIR train resize 192 +python -m data.inat preprocess $DATA_DIR val resize 256 +python -m data.inat preprocess $DATA_DIR train resize 256 ``` 5. Login to Wandb diff --git a/config.py b/config.py index 6a86ccbc..f8cf7fdc 100644 --- a/config.py +++ b/config.py @@ -19,8 +19,6 @@ # Data settings # ----------------------------------------------------------------------------- _C.DATA = CN() -# Batch size for a single GPU, could be overwritten by command line argument -_C.DATA.BATCH_SIZE = 128 # Path to dataset, could be overwritten by command line argument _C.DATA.DATA_PATH = "" # Dataset name @@ -146,6 +144,10 @@ # Training settings # ----------------------------------------------------------------------------- _C.TRAIN = CN() +# Batch size for a single GPU, could be overwritten by command line argument +_C.TRAIN.DEVICE_BATCH_SIZE = 128 +# Global batch size = DEVICE_BATCH_SIZE * N_PROCS * ACCUMULATION_STEPS +_C.TRAIN.GLOBAL_BATCH_SIZE = 1024 _C.TRAIN.START_EPOCH = 0 _C.TRAIN.EPOCHS = 300 _C.TRAIN.WARMUP_EPOCHS = 20 @@ -238,6 +240,16 @@ _C.TEST.SEQUENTIAL = False _C.TEST.SHUFFLE = False +# ----------------------------------------------------------------------------- +# Experiment Settings +# ----------------------------------------------------------------------------- +_C.EXPERIMENT = CN() +# The experiment name. This is a human-readable name that is easy to read. +_C.EXPERIMENT.NAME = "default-dragonfruit" +# The wandb id for logging. +# Generate this id with scripts/generate_wandb_id +_C.EXPERIMENT.WANDB_ID = "" + # ----------------------------------------------------------------------------- # Misc # ----------------------------------------------------------------------------- @@ -254,8 +266,6 @@ _C.AMP_OPT_LEVEL = "" # Path to output folder, overwritten by command line argument _C.OUTPUT = "" -# Tag of experiment, overwritten by command line argument -_C.TAG = "default" # Frequency to save checkpoint _C.SAVE_FREQ = 1 # Frequency to logging info @@ -302,9 +312,9 @@ def _check_args(name): # merge from specific arguments if _check_args("batch_size"): - config.DATA.BATCH_SIZE = args.batch_size + config.TRAIN.DEVICE_BATCH_SIZE = args.batch_size if _check_args("data_path"): - config.DATA.DATA_PATH = args.data_path + config.DATA.DATA_PATH = os.path.abspath(args.data_path) if _check_args("zip"): config.DATA.ZIP_MODE = True if _check_args("cache_mode"): @@ -313,8 +323,6 @@ def _check_args(name): config.MODEL.PRETRAINED = args.pretrained if _check_args("resume"): config.MODEL.RESUME = args.resume - if _check_args("accumulation_steps"): - config.TRAIN.ACCUMULATION_STEPS = args.accumulation_steps if _check_args("use_checkpoint"): config.TRAIN.USE_CHECKPOINT = True if _check_args("amp_opt_level"): @@ -325,8 +333,6 @@ def _check_args(name): config.AMP_ENABLE = False if _check_args("output"): config.OUTPUT = args.output - if _check_args("tag"): - config.TAG = args.tag if _check_args("eval"): config.EVAL_MODE = True if _check_args("throughput"): @@ -350,8 +356,38 @@ def _check_args(name): # set local rank for distributed training config.LOCAL_RANK = int(os.environ["LOCAL_RANK"]) + # Use this to calculate accumulation steps + if "LOCAL_WORLD_SIZE" in os.environ: + + def divide_cleanly(a, b): + assert a % b == 0, f"{a} / {b} has remainder {a % b}" + return a // b + + n_procs = int(os.environ["LOCAL_WORLD_SIZE"]) + desired_device_batch_size = divide_cleanly( + config.TRAIN.GLOBAL_BATCH_SIZE, n_procs + ) + actual_device_batch_size = config.TRAIN.DEVICE_BATCH_SIZE + + if actual_device_batch_size > desired_device_batch_size: + print( + f"Decreasing device batch size from {actual_device_batch_size} to {desired_device_batch_size} so your global bath size is {config.TRAIN.GLOBAL_BATCH_SIZE}, not {desired_device_batch_size * n_procs}!" + ) + config.TRAIN.ACCUMULATION_STEPS = 1 + config.TRAIN.DEVICE_BATCH_SIZE = desired_device_batch_size + elif desired_device_batch_size == actual_device_batch_size: + config.TRAIN.ACCUMULATION_STEPS = 1 + else: + assert desired_device_batch_size > actual_device_batch_size + config.TRAIN.ACCUMULATION_STEPS = divide_cleanly( + desired_device_batch_size, actual_device_batch_size + ) + print( + f"Using {config.TRAIN.ACCUMULATION_STEPS} accumulation steps so your global batch size is {config.TRAIN.GLOBAL_BATCH_SIZE}, not {actual_device_batch_size * n_procs}!" + ) + # output folder - config.OUTPUT = os.path.join(config.OUTPUT, config.MODEL.NAME, config.TAG) + config.OUTPUT = os.path.join(config.OUTPUT, config.EXPERIMENT.NAME) config.freeze() diff --git a/configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml b/configs/hierarchical-vision-project/funky-banana-192.yaml similarity index 85% rename from configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml rename to configs/hierarchical-vision-project/funky-banana-192.yaml index 63f05fac..a03a065e 100644 --- a/configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml +++ b/configs/hierarchical-vision-project/funky-banana-192.yaml @@ -1,12 +1,10 @@ DATA: DATASET: inat21 IMG_SIZE: 192 - DATA_PATH: /research/nfs_su_809/cv_datasets/inat21/train_val_192 NUM_WORKERS: 32 - BATCH_SIZE: 256 MODEL: TYPE: swinv2 - NAME: funky-banana-192 + NAME: swinv2_base DROP_PATH_RATE: 0.2 SWINV2: EMBED_DIM: 128 @@ -16,7 +14,7 @@ MODEL: TRAIN: # Want a global batch size of 2048 because SwinV2 was trained on 16 V100s with batch size 128 (I think) # But we are going to use a global batch size of 1024 because it's faster (throughput). - ACCUMULATION_STEPS: 1 + GLOBAL_BATCH_SIZE: 1024 # We are using limited epochs based on pre-training configs for imagenet22k # Then we will pre-train on 256x256 for 30 epochs @@ -25,3 +23,7 @@ TRAIN: WEIGHT_DECAY: 0.1 SAVE_FREQ: 4 + +EXPERIMENT: + NAME: funky-banana-192 + WANDB_ID: 2y91axvs diff --git a/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml b/configs/hierarchical-vision-project/groovy-grape-256.yaml similarity index 72% rename from configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml rename to configs/hierarchical-vision-project/groovy-grape-256.yaml index 2e88e62b..d000dccd 100644 --- a/configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21_hierarchical_lr1.25_ft.yaml +++ b/configs/hierarchical-vision-project/groovy-grape-256.yaml @@ -1,13 +1,10 @@ DATA: DATASET: inat21 - NUM_WORKERS: 32 - BATCH_SIZE: 16 IMG_SIZE: 256 - DATA_PATH: /mnt/10tb/data/inat21/resize-256 + NUM_WORKERS: 32 MODEL: TYPE: swinv2 - PRETRAINED: /mnt/10tb/models/swinv2_base_patch4_window12_192_inat21_hierarchical_lr1.25_v0_epoch_89.pth - NAME: groovy-grape-256 + NAME: swinv2_base_window16 DROP_PATH_RATE: 0.2 SWINV2: EMBED_DIM: 128 @@ -16,8 +13,7 @@ MODEL: WINDOW_SIZE: 16 PRETRAINED_WINDOW_SIZES: [ 12, 12, 12, 6 ] TRAIN: - # Global batch size of 1024 - ACCUMULATION_STEPS: 8 + GLOBAL_BATCH_SIZE: 1024 EPOCHS: 30 WARMUP_EPOCHS: 5 @@ -31,4 +27,8 @@ TRAIN: HIERARCHICAL_COEFFS: [ 8, 5.65, 4, 2.82, 2, 1.41, 1 ] +EXPERIMENT: + NAME: groovy-grape-256 + WANDB_ID: 11o47wpm + HIERARCHICAL: true diff --git a/data/build.py b/data/build.py index c9b140cd..89dd554b 100644 --- a/data/build.py +++ b/data/build.py @@ -58,7 +58,7 @@ def build_loader(config): # Check if we are overfitting some subset of the training data for debugging if config.TRAIN.OVERFIT_BATCHES > 0: - n_examples = config.TRAIN.OVERFIT_BATCHES * config.DATA.BATCH_SIZE + n_examples = config.TRAIN.OVERFIT_BATCHES * config.TRAIN.DEVICE_BATCH_SIZE indices = random.sample(range(len(dataset_train)), n_examples) dataset_train = Subset(dataset_train, indices) @@ -82,7 +82,7 @@ def build_loader(config): data_loader_train = torch.utils.data.DataLoader( dataset_train, sampler=sampler_train, - batch_size=config.DATA.BATCH_SIZE, + batch_size=config.TRAIN.DEVICE_BATCH_SIZE, num_workers=config.DATA.NUM_WORKERS, pin_memory=config.DATA.PIN_MEMORY, drop_last=True, @@ -91,7 +91,7 @@ def build_loader(config): data_loader_val = torch.utils.data.DataLoader( dataset_val, sampler=sampler_val, - batch_size=config.DATA.BATCH_SIZE, + batch_size=config.TRAIN.DEVICE_BATCH_SIZE, shuffle=False, num_workers=config.DATA.NUM_WORKERS, pin_memory=config.DATA.PIN_MEMORY, @@ -164,7 +164,7 @@ def build_dataset(is_train, config): dataset = datasets.ImageFolder(root, transform=transform) nb_classes = 10_000 else: - raise NotImplementedError("We only support ImageNet Now.") + raise NotImplementedError("We only support ImageNet now.") return dataset, nb_classes @@ -220,7 +220,7 @@ def build_transform(is_train, config): mean, std = data_mean_std[config.DATA.DATASET] else: raise RuntimeError( - f"Can't find mean/std for {config.DATA.DATASET} at {config.DATA.DATASET}. Please add it to data/constants.py (try using python -m data.inat normalize for iNat)." + f"Can't find mean/std for {config.DATA.DATASET} at {config.DATA.DATA_PATH}. Please add it to data/constants.py (try using python -m data.inat normalize for iNat)." ) t.append(transforms.ToTensor()) diff --git a/data/constants.py b/data/constants.py index 001823ae..c7e362cb 100644 --- a/data/constants.py +++ b/data/constants.py @@ -6,7 +6,7 @@ torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), ), - "/research/nfs_su_809/cv_datasets/inat21/train_val_192": ( + "/local/scratch/cv_datasets/inat21/train_val_192": ( torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), ), @@ -18,5 +18,9 @@ torch.tensor([0.23768986761569977, 0.22925858199596405, 0.2476460039615631]), torch.tensor([0.4632672071456909, 0.480050653219223, 0.37618669867515564]), ), + "/local/scratch/cv_datasets/inat21/resize-256": ( + torch.tensor([0.23768986761569977, 0.22925858199596405, 0.2476460039615631]), + torch.tensor([0.4632672071456909, 0.480050653219223, 0.37618669867515564]), + ), "imagenet": (IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD), } diff --git a/docs/experiments/funky-banana.md b/docs/experiments/funky-banana.md index d5b88fde..c50c9b13 100644 --- a/docs/experiments/funky-banana.md +++ b/docs/experiments/funky-banana.md @@ -5,7 +5,7 @@ We train for 90 epochs at 192x192, then tune for 30 epochs at 256x256. ```yaml configs: -- configs/swinv2/swinv2_base_patch4_window12_192_inat21.yaml +- configs/hierarchical-vision-project/funky-banana-192.yaml - configs/swinv2/swinv2_base_patch4_window12to16_192to256_inat21.yaml codename: funky-banana ``` diff --git a/logger.py b/logger.py index ff3bf7c9..ab1d1aef 100644 --- a/logger.py +++ b/logger.py @@ -11,10 +11,11 @@ import sys import torch -import wandb from termcolor import colored from torch.utils.tensorboard import SummaryWriter +import wandb + @functools.lru_cache() def create_logger(output_dir, dist_rank=0, name=""): @@ -90,23 +91,24 @@ def init(self, config): if self.rank != 0: return - wandb.init( + kwargs = dict( config=config, project="hierarchical-vision", - dir="./runs", - resume=bool(config.MODEL.RESUME), - name=config.MODEL.NAME, + name=config.EXPERIMENT.NAME, ) + + if not config.EXPERIMENT.WANDB_ID: + print("Cannot resume wandb run because no id was provided!") + else: + kwargs["id"] = config.EXPERIMENT.WANDB_ID + kwargs["resume"] = "allow" + + wandb.init(**kwargs) + # Validation metrics - wandb.define_metric( - "val/loss", step_metric="epoch", summary="best", objective="max" - ) - wandb.define_metric( - "val/acc1", step_metric="epoch", summary="best", objective="max" - ) - wandb.define_metric( - "val/acc5", step_metric="epoch", summary="best", objective="max" - ) + wandb.define_metric("val/loss", step_metric="epoch", summary="max") + wandb.define_metric("val/acc1", step_metric="epoch", summary="max") + wandb.define_metric("val/acc5", step_metric="epoch", summary="max") # Training metrics wandb.define_metric("train/batch_time", step_metric="step", summary="last") diff --git a/main.py b/main.py index b6311cf2..725b119a 100644 --- a/main.py +++ b/main.py @@ -33,7 +33,7 @@ from utils import ( NativeScalerWithGradNormCount, auto_resume_helper, - batch_size, + batch_size_of, load_checkpoint, load_pretrained, reduce_tensor, @@ -334,7 +334,7 @@ def train_one_epoch( # We divide by accumulation steps (not sure why) but it makes # the logged values look weird. So I multiply by it to fix that. loss_meter.update( - loss.item() * config.TRAIN.ACCUMULATION_STEPS, batch_size(targets) + loss.item() * config.TRAIN.ACCUMULATION_STEPS, batch_size_of(targets) ) if grad_norm is not None: # loss_scaler return None if not update @@ -497,13 +497,22 @@ def throughput(data_loader, model, logger): # linear scale the learning rate according to total batch size, may not be optimal linear_scaled_lr = ( - config.TRAIN.BASE_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + config.TRAIN.BASE_LR + * config.TRAIN.DEVICE_BATCH_SIZE + * dist.get_world_size() + / 512.0 ) linear_scaled_warmup_lr = ( - config.TRAIN.WARMUP_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + config.TRAIN.WARMUP_LR + * config.TRAIN.DEVICE_BATCH_SIZE + * dist.get_world_size() + / 512.0 ) linear_scaled_min_lr = ( - config.TRAIN.MIN_LR * config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0 + config.TRAIN.MIN_LR + * config.TRAIN.DEVICE_BATCH_SIZE + * dist.get_world_size() + / 512.0 ) # gradient accumulation also need to scale the learning rate if config.TRAIN.ACCUMULATION_STEPS > 1: @@ -520,7 +529,9 @@ def throughput(data_loader, model, logger): os.makedirs(config.OUTPUT, exist_ok=True) logger = create_logger( - output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" + output_dir=config.OUTPUT, + dist_rank=dist.get_rank(), + name=f"{config.EXPERIMENT.NAME}", ) wandb_writer = WandbWriter(rank=dist.get_rank()) wandb_writer.init(config) diff --git a/main_moe.py b/main_moe.py index 082d9c59..9a611cfa 100644 --- a/main_moe.py +++ b/main_moe.py @@ -499,7 +499,9 @@ def throughput(data_loader, model, logger): os.makedirs(config.OUTPUT, exist_ok=True) logger = create_logger( - output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" + output_dir=config.OUTPUT, + dist_rank=dist.get_rank(), + name=f"{config.EXPERIMENT.NAME}", ) if dist.get_rank() == 0: diff --git a/main_simmim_ft.py b/main_simmim_ft.py index 2f0df1d7..e0ea09ad 100644 --- a/main_simmim_ft.py +++ b/main_simmim_ft.py @@ -421,7 +421,9 @@ def throughput(data_loader, model, logger): os.makedirs(config.OUTPUT, exist_ok=True) logger = create_logger( - output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" + output_dir=config.OUTPUT, + dist_rank=dist.get_rank(), + name=f"{config.EXPERIMENT.NAME}", ) if dist.get_rank() == 0: diff --git a/main_simmim_pt.py b/main_simmim_pt.py index b4d91dad..decca083 100644 --- a/main_simmim_pt.py +++ b/main_simmim_pt.py @@ -277,7 +277,9 @@ def train_one_epoch(config, model, data_loader, optimizer, epoch, lr_scheduler, os.makedirs(config.OUTPUT, exist_ok=True) logger = create_logger( - output_dir=config.OUTPUT, dist_rank=dist.get_rank(), name=f"{config.MODEL.NAME}" + output_dir=config.OUTPUT, + dist_rank=dist.get_rank(), + name=f"{config.EXPERIMENT.NAME}", ) if dist.get_rank() == 0: diff --git a/scripts/generate_wandb_id.py b/scripts/generate_wandb_id.py new file mode 100644 index 00000000..b32fc512 --- /dev/null +++ b/scripts/generate_wandb_id.py @@ -0,0 +1,3 @@ +import wandb.util + +print(wandb.util.generate_id()) diff --git a/scripts/train.fish b/scripts/train.fish index a0c1c351..a02a7905 100644 --- a/scripts/train.fish +++ b/scripts/train.fish @@ -6,59 +6,97 @@ function usage echo echo "Arguments:" echo " -h/--help print this message" - echo " --config which config file to use (should be YAML)" - echo " --debug run with only one process rather than 8 so pdb is useable." - echo " --nprocs number of processes" - echo " --tag tag for the run (I use v0, v1, etc)" - echo " --venv path to the virtual environment (default ./venv/)" + echo + echo " --batch-size GPU batch size (default 256)" + echo " --config which config file to use (required)" + echo " --data path to dataset (required)" + echo " --debug run with only one process rather than 8 so pdb is useable." + echo " --port the master processes port (default 12345)" + echo " --nprocs number of processes" end +# Parse arguments set -l options (fish_opt --short h --long help) +set -a options (fish_opt --short b --long batch-size --required-val --long-only) set -a options (fish_opt --short c --long config --required-val --long-only) -set -a options (fish_opt --short t --long tag --required-val --long-only) -set -a options (fish_opt --short p --long nprocs --required-val --long-only) -set -a options (fish_opt --short v --long venv --required-val --long-only) -set -a options (fish_opt --short d --long debug --long-only) +# --short g for debuG because --data uses --short d +set -a options (fish_opt --short g --long debug --long-only) +set -a options (fish_opt --short d --long data --required-val --long-only) +set -a options (fish_opt --short p --long port --required-val --long-only) +set -a options (fish_opt --short n --long nprocs --required-val --long-only) argparse $options -- $argv + +# Print help if not test -z $_flag_help usage exit 0 end + +# Check environment variables +if not string length -q -- $VENV + echo "You need to provide a VENV environment variable." + echo "Try:" + echo " source scripts/local.fish" + exit 2 +end + +if not string length -q -- $RUN_OUTPUT + echo "You need to provide a RUN_OUTPUT environment variable." + echo "Try:" + echo " source scripts/local.fish" + exit 2 +end + +set launcher $VENV/bin/torchrun + +# --batch size +set batch_size 256 +if not test -z $_flag_batch_size + set batch_size $_flag_batch_size +end + +# --config if test -z $_flag_config echo "You must provide a --config argument!" exit 1 end -if test -z $_flag_tag - echo "You must provide a --tag argument!" +# --config +if test -z $_flag_data + echo "You must provide a --data argument!" exit 1 end -set venv ./venv -if not test -z $_flag_venv - set venv $_flag_venv +# --port +set port 12345 +if not test -z $_flag_port + set port $_flag_port end +# --nprocs set nprocs 8 if not test -z $_flag_nprocs set nprocs $_flag_nprocs end -set launcher $venv/bin/torchrun -set launcher_args --nproc_per_node $nprocs - +# --debug +# Has to come after nprocs because it modifies nprocs if not test -z $_flag_debug - set launcher_args --nproc_per_node 1 + set nprocs 1 end -$launcher $launcher_args \ - --master_port 12345 \ +# echo $launcher $nprocs $port $_flag_config $RUN_OUTPUT $batch_size +# exit 1 + +OMP_NUM_THREADS=32 $launcher --nproc_per_node $nprocs \ + --master_port $port \ main.py \ --cfg $_flag_config \ - --output runs \ - --tag $_flag_tag \ + --data-path $_flag_data \ + --output $RUN_OUTPUT \ + --batch-size $batch_size \ --fused_window_process \ --fused_layernorm diff --git a/utils.py b/utils.py index f7a756aa..93a16318 100644 --- a/utils.py +++ b/utils.py @@ -378,7 +378,7 @@ def ampscaler_get_grad_norm(parameters, norm_type: float = 2.0) -> torch.Tensor: return total_norm -def batch_size(tensor_or_list): +def batch_size_of(tensor_or_list): if isinstance(tensor_or_list, torch.Tensor): return tensor_or_list.size(0) elif isinstance(tensor_or_list, list): From 615b9bd2952b22b07adcf5b8be282b97233254f4 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Wed, 2 Nov 2022 14:52:59 +0000 Subject: [PATCH 40/47] Update fuzzy-fig-192 --- README.md | 15 +++++++++ .../fuzzy-fig-192.yaml | 32 +++++++++++++++++++ scripts/train.fish | 6 ++-- 3 files changed, 50 insertions(+), 3 deletions(-) create mode 100644 configs/hierarchical-vision-project/fuzzy-fig-192.yaml diff --git a/README.md b/README.md index 065e4bb9..e333743d 100644 --- a/README.md +++ b/README.md @@ -82,6 +82,21 @@ python -m data.inat preprocess $DATA_DIR train resize 256 wandb login ``` +6. Set up an `env.fish` file: + +You need to provide `$VENV` and a `$RUN_OUTPUT` environment variables. +I recommend using a file to save these variables. + +In fish: + +```fish +# scripts/env.fish +set -gx VENV venv +set -gx RUN_OUTPUT /mnt/10tb/models/hierarchical-vision +``` + +Then run `source scripts/env.fish` + ## AWS Helpers Uninstall v1 of awscli: diff --git a/configs/hierarchical-vision-project/fuzzy-fig-192.yaml b/configs/hierarchical-vision-project/fuzzy-fig-192.yaml new file mode 100644 index 00000000..d68232a5 --- /dev/null +++ b/configs/hierarchical-vision-project/fuzzy-fig-192.yaml @@ -0,0 +1,32 @@ +DATA: + DATASET: inat21 + IMG_SIZE: 192 + NUM_WORKERS: 32 +MODEL: + TYPE: swinv2 + NAME: swinv2_base_window12 + DROP_PATH_RATE: 0.2 + SWINV2: + EMBED_DIM: 128 + DEPTHS: [ 2, 2, 18, 2 ] + NUM_HEADS: [ 4, 8, 16, 32 ] + WINDOW_SIZE: 12 +TRAIN: + # Want a global batch size of 2048 because SwinV2 was trained on 16 V100s with batch size 128 (I think) + # But we are going to use a global batch size of 1024 because it's faster (throughput). + GLOBAL_BATCH_SIZE: 1024 + + # We are using limited epochs based on pre-training configs for imagenet22k + # Then we will pre-train on 256x256 for 30 epochs + EPOCHS: 90 + WARMUP_EPOCHS: 5 + WEIGHT_DECAY: 0.1 + + # Use 1/4 of original learning rates + BASE_LR: 1.25e-4 + WARMUP_LR: 1.25e-7 + MIN_LR: 1.25e-6 + +EXPERIMENT: + NAME: fuzzy-fig-192 + WANDB_ID: 2c0bq1h2 diff --git a/scripts/train.fish b/scripts/train.fish index a02a7905..dd5bc6b6 100644 --- a/scripts/train.fish +++ b/scripts/train.fish @@ -12,7 +12,7 @@ function usage echo " --data path to dataset (required)" echo " --debug run with only one process rather than 8 so pdb is useable." echo " --port the master processes port (default 12345)" - echo " --nprocs number of processes" + echo " --nprocs number of processes (default 8)" end # Parse arguments @@ -39,14 +39,14 @@ end if not string length -q -- $VENV echo "You need to provide a VENV environment variable." echo "Try:" - echo " source scripts/local.fish" + echo " source scripts/env.fish" exit 2 end if not string length -q -- $RUN_OUTPUT echo "You need to provide a RUN_OUTPUT environment variable." echo "Try:" - echo " source scripts/local.fish" + echo " source scripts/env.fish" exit 2 end From 3c663b14ddb57a058e2b6f7753da90e7258c2bd2 Mon Sep 17 00:00:00 2001 From: Samuel Stevens Date: Wed, 2 Nov 2022 11:04:11 -0400 Subject: [PATCH 41/47] Improved options for evaluating pre-trained models/checkpoints --- data/constants.py | 4 ++++ main.py | 1 + scripts/train.fish | 24 ++++++++++++++++++++---- utils.py | 13 ++++++++----- 4 files changed, 33 insertions(+), 9 deletions(-) diff --git a/data/constants.py b/data/constants.py index c7e362cb..4cab412f 100644 --- a/data/constants.py +++ b/data/constants.py @@ -6,6 +6,10 @@ torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), ), + "/local/scratch/cv_datasets/inat21/resize-192": ( + torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), + torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), + ), "/local/scratch/cv_datasets/inat21/train_val_192": ( torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), diff --git a/main.py b/main.py index 725b119a..2e6c244e 100644 --- a/main.py +++ b/main.py @@ -217,6 +217,7 @@ def main(config): logger.info( f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%" ) + logger.info("Previously reported best accuracy: %.2f", max_accuracy) if config.EVAL_MODE: return diff --git a/scripts/train.fish b/scripts/train.fish index dd5bc6b6..dacafd4c 100644 --- a/scripts/train.fish +++ b/scripts/train.fish @@ -25,7 +25,7 @@ set -a options (fish_opt --short d --long data --required-val --long-only) set -a options (fish_opt --short p --long port --required-val --long-only) set -a options (fish_opt --short n --long nprocs --required-val --long-only) -argparse $options -- $argv +argparse --ignore-unknown $options -- $argv # Print help @@ -88,8 +88,23 @@ if not test -z $_flag_debug set nprocs 1 end -# echo $launcher $nprocs $port $_flag_config $RUN_OUTPUT $batch_size -# exit 1 +if not set -q var[1] + # There is at least one remaining variable. + # Confirm that this is really what you want to do + echo "Do you want to pass the additional arguments" + echo + echo " $argv" + echo + echo "to main.py?" + read --local --prompt-str "[y/N] " confirm + switch $confirm + case Y y + echo "Okay!" + case '*' + echo "Exiting..." + exit 0 + end +end OMP_NUM_THREADS=32 $launcher --nproc_per_node $nprocs \ --master_port $port \ @@ -99,4 +114,5 @@ OMP_NUM_THREADS=32 $launcher --nproc_per_node $nprocs \ --output $RUN_OUTPUT \ --batch-size $batch_size \ --fused_window_process \ - --fused_layernorm + --fused_layernorm \ + $argv # include additional arguments diff --git a/utils.py b/utils.py index 93a16318..8e36f973 100644 --- a/utils.py +++ b/utils.py @@ -25,6 +25,14 @@ def load_checkpoint(config, model, optimizer, lr_scheduler, loss_scaler, logger) msg = model.load_state_dict(checkpoint["model"], strict=False) logger.info(msg) max_accuracy = 0.0 + + if "max_accuracy" in checkpoint: + max_accuracy = checkpoint["max_accuracy"] + + config.defrost() + config.TRAIN.START_EPOCH = checkpoint["epoch"] + 1 + config.freeze() + if ( not config.EVAL_MODE and "optimizer" in checkpoint @@ -33,16 +41,11 @@ def load_checkpoint(config, model, optimizer, lr_scheduler, loss_scaler, logger) ): optimizer.load_state_dict(checkpoint["optimizer"]) lr_scheduler.load_state_dict(checkpoint["lr_scheduler"]) - config.defrost() - config.TRAIN.START_EPOCH = checkpoint["epoch"] + 1 - config.freeze() if "scaler" in checkpoint: loss_scaler.load_state_dict(checkpoint["scaler"]) logger.info( f"=> loaded successfully '{config.MODEL.RESUME}' (epoch {checkpoint['epoch']})" ) - if "max_accuracy" in checkpoint: - max_accuracy = checkpoint["max_accuracy"] del checkpoint torch.cuda.empty_cache() From 1f47f00fb499660c7cb5f82d4d13606e1cb54684 Mon Sep 17 00:00:00 2001 From: Sam Stevens Date: Wed, 2 Nov 2022 15:07:30 +0000 Subject: [PATCH 42/47] Rebasing --- .../funky-banana-192.yaml | 4 +- data/constants.py | 6 +- scripts/parse_logs.py | 198 +++++++++++------- utils.py | 2 +- 4 files changed, 124 insertions(+), 86 deletions(-) diff --git a/configs/hierarchical-vision-project/funky-banana-192.yaml b/configs/hierarchical-vision-project/funky-banana-192.yaml index a03a065e..61b72d1a 100644 --- a/configs/hierarchical-vision-project/funky-banana-192.yaml +++ b/configs/hierarchical-vision-project/funky-banana-192.yaml @@ -4,7 +4,7 @@ DATA: NUM_WORKERS: 32 MODEL: TYPE: swinv2 - NAME: swinv2_base + NAME: swinv2_base_window12 DROP_PATH_RATE: 0.2 SWINV2: EMBED_DIM: 128 @@ -26,4 +26,4 @@ SAVE_FREQ: 4 EXPERIMENT: NAME: funky-banana-192 - WANDB_ID: 2y91axvs + WANDB_ID: 2a0tgwo2 diff --git a/data/constants.py b/data/constants.py index 4cab412f..8d292fa9 100644 --- a/data/constants.py +++ b/data/constants.py @@ -3,14 +3,10 @@ data_mean_std = { "/mnt/10tb/data/inat21/resize-192": ( - torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), - ), - "/local/scratch/cv_datasets/inat21/resize-192": ( torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), - torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), ), - "/local/scratch/cv_datasets/inat21/train_val_192": ( + "/local/scratch/cv_datasets/inat21/resize-192": ( torch.tensor([0.23754851520061493, 0.22912880778312683, 0.24746596813201904]), torch.tensor([0.4632684290409088, 0.48004600405693054, 0.37628623843193054]), ), diff --git a/scripts/parse_logs.py b/scripts/parse_logs.py index 5815c863..73f0c59d 100644 --- a/scripts/parse_logs.py +++ b/scripts/parse_logs.py @@ -16,6 +16,7 @@ def parse_args(): parser.add_argument( "--last", help="How many of the latest epochs to look at.", default=10, type=int ) + parser.add_argument("--name", help="Run name") return parser.parse_args() @@ -32,6 +33,39 @@ class ValidationLine: acc5: float mean_acc5: float + pattern0 = r""" + ^\[.*?\]\ + INFO:\ Test:\ \ + \[\ ?(?P\d+)/(?P\d+)\] + \ \ + eta:\ \d+:\d\d:\d\d + \ \ + loss:\ (?P[\w.]+)\ \((?P[\w.]+)\) + \ \ + acc1:\ (?P[\w.]+)\ \((?P[\w.]+)\) + \ \ + acc5:\ (?P[\w.]+)\ \((?P[\w.]+)\) + $ + """ + + pattern1 = r""" + ^\[.*?\]\ + \(main.py\ \d+\):\ + INFO\ Test:\ + \[(?P\d+)/(?P\d+)\] + \t + Time\ \d+.\d+\ \(\d+.\d+\) + \t + Loss\ (?P[\w.]+)\ \((?P[\w.]+)\) + \t + Acc@1\ (?P[\w.]+)\ \((?P[\w.]+)\) + \t + Acc@5\ (?P[\w.]+)\ \((?P[\w.]+)\) + \t + Mem\ (?P.*) + $ + """ + @classmethod def from_raw_line(cls, line, last_train): if "Test" not in line: @@ -40,28 +74,12 @@ def from_raw_line(cls, line, last_train): # Example line: # [2022-06-08 07:34:50 swinv2_large_patch4_window7_224_inat21](main.py 258): INFO Test: [0/196] Time 1.895 (1.895) Loss 0.8066 (0.8066) Acc@1 84.766 (84.766) Acc@5 95.312 (95.312) Mem 36916MB - pattern = r""" -^\[.*?\]\ -\(main.py\ \d+\):\ -INFO\ Test:\ -\[(?P\d+)/(?P\d+)\] -\t -Time\ \d+.\d+\ \(\d+.\d+\) -\t -Loss\ (?P[\w.]+)\ \((?P[\w.]+)\) -\t -Acc@1\ (?P[\w.]+)\ \((?P[\w.]+)\) -\t -Acc@5\ (?P[\w.]+)\ \((?P[\w.]+)\) -\t -Mem\ (?P.*) -$ -""" - - match = re.match(pattern, line, re.VERBOSE) + match = re.match(cls.pattern0, line, re.VERBOSE) or re.match( + cls.pattern1, line, re.VERBOSE + ) if not match: - print(repr(line)) + print("Couldn't match validation line:", repr(line)) return None epoch = 0 @@ -96,58 +114,71 @@ class TrainLine: loss_scale: float mean_loss_scale: float + pattern0 = r""" + ^\[.*?\]\ # [2022-06-08 08:35:04 ...] + INFO:\ Epoch:\ \[(?P\d+)\] + \ \ + \[\ *(?P\d+)/(?P\d+)\] # [700/5247] + \ \ # + eta:\ (\d\ day,\ )?\d+:\d\d:\d\d # eta 3:11:57 + \ \ # + lr:\ (?P\d\.\d+) # lr 0.000040 + \ \ # + loss:\ (?P[\w.]+)\ \((?P[\w.]+)\) # loss 4.3632 (3.7333) + \ \ # + grad_norm:\ (?P[\w.]+)\ \((?P[\w.]+)\) # grad_norm 9.5996 (inf) + \ \ + time:\ (?P