update

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

app.py

@@ -1,7 +1,414 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+import argparse
+import datetime
+import numpy as np
+import time
+import torch
+import torch.backends.cudnn as cudnn
+import json
+
+from pathlib import Path
+from utils import get_root_logger
+from timm.models import create_model
+import models_v2
+import requests
+
+import utils
+import time
+import sys
+import datetime
+import os
+from snnet import SNNet, SNNetv2
+import warnings
+
+warnings.filterwarnings("ignore")
+from fvcore.nn import FlopCountAnalysis
+
+from PIL import Image
 import gradio as gr
+import plotly.express as px
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('DeiT training and evaluation script', add_help=False)
+    parser.add_argument('--batch-size', default=64, type=int)
+    parser.add_argument('--epochs', default=300, type=int)
+    parser.add_argument('--bce-loss', action='store_true')
+    parser.add_argument('--unscale-lr', action='store_true')
+
+    # Model parameters
+    parser.add_argument('--model', default='deit_base_patch16_224', type=str, metavar='MODEL',
+                        help='Name of model to train')
+    parser.add_argument('--input-size', default=224, type=int, help='images input size')
+
+    parser.add_argument('--drop', type=float, default=0.0, metavar='PCT',
+                        help='Dropout rate (default: 0.)')
+    parser.add_argument('--drop-path', type=float, default=0.1, metavar='PCT',
+                        help='Drop path rate (default: 0.1)')
+
+    parser.add_argument('--model-ema', action='store_true')
+    parser.add_argument('--no-model-ema', action='store_false', dest='model_ema')
+    parser.set_defaults(model_ema=True)
+    parser.add_argument('--model-ema-decay', type=float, default=0.99996, help='')
+    parser.add_argument('--model-ema-force-cpu', action='store_true', default=False, help='')
+
+    # Optimizer parameters
+    parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER',
+                        help='Optimizer (default: "adamw"')
+    parser.add_argument('--opt-eps', default=1e-8, type=float, metavar='EPSILON',
+                        help='Optimizer Epsilon (default: 1e-8)')
+    parser.add_argument('--opt-betas', default=None, type=float, nargs='+', metavar='BETA',
+                        help='Optimizer Betas (default: None, use opt default)')
+    parser.add_argument('--clip-grad', type=float, default=None, metavar='NORM',
+                        help='Clip gradient norm (default: None, no clipping)')
+    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
+                        help='SGD momentum (default: 0.9)')
+    parser.add_argument('--weight-decay', type=float, default=0.05,
+                        help='weight decay (default: 0.05)')
+    # Learning rate schedule parameters
+    parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER',
+                        help='LR scheduler (default: "cosine"')
+    parser.add_argument('--lr', type=float, default=5e-4, metavar='LR',
+                        help='learning rate (default: 5e-4)')
+    parser.add_argument('--lr-noise', type=float, nargs='+', default=None, metavar='pct, pct',
+                        help='learning rate noise on/off epoch percentages')
+    parser.add_argument('--lr-noise-pct', type=float, default=0.67, metavar='PERCENT',
+                        help='learning rate noise limit percent (default: 0.67)')
+    parser.add_argument('--lr-noise-std', type=float, default=1.0, metavar='STDDEV',
+                        help='learning rate noise std-dev (default: 1.0)')
+    parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR',
+                        help='warmup learning rate (default: 1e-6)')
+    parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',
+                        help='lower lr bound for cyclic schedulers that hit 0 (1e-5)')
+
+    parser.add_argument('--decay-epochs', type=float, default=30, metavar='N',
+                        help='epoch interval to decay LR')
+    parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',
+                        help='epochs to warmup LR, if scheduler supports')
+    parser.add_argument('--cooldown-epochs', type=int, default=10, metavar='N',
+                        help='epochs to cooldown LR at min_lr, after cyclic schedule ends')
+    parser.add_argument('--patience-epochs', type=int, default=10, metavar='N',
+                        help='patience epochs for Plateau LR scheduler (default: 10')
+    parser.add_argument('--decay-rate', '--dr', type=float, default=0.1, metavar='RATE',
+                        help='LR decay rate (default: 0.1)')
+
+    # Augmentation parameters
+    parser.add_argument('--color-jitter', type=float, default=0.3, metavar='PCT',
+                        help='Color jitter factor (default: 0.3)')
+    parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
+                        help='Use AutoAugment policy. "v0" or "original". " + \
+                             "(default: rand-m9-mstd0.5-inc1)'),
+    parser.add_argument('--smoothing', type=float, default=0.1, help='Label smoothing (default: 0.1)')
+    parser.add_argument('--train-interpolation', type=str, default='bicubic',
+                        help='Training interpolation (random, bilinear, bicubic default: "bicubic")')
+
+    parser.add_argument('--repeated-aug', action='store_true')
+    parser.add_argument('--no-repeated-aug', action='store_false', dest='repeated_aug')
+    parser.set_defaults(repeated_aug=True)
+
+    parser.add_argument('--train-mode', action='store_true')
+    parser.add_argument('--no-train-mode', action='store_false', dest='train_mode')
+    parser.set_defaults(train_mode=True)
+
+    parser.add_argument('--ThreeAugment', action='store_true')  # 3augment
+
+    parser.add_argument('--src', action='store_true')  # simple random crop
+
+    # * Random Erase params
+    parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
+                        help='Random erase prob (default: 0.25)')
+    parser.add_argument('--remode', type=str, default='pixel',
+                        help='Random erase mode (default: "pixel")')
+    parser.add_argument('--recount', type=int, default=1,
+                        help='Random erase count (default: 1)')
+    parser.add_argument('--resplit', action='store_true', default=False,
+                        help='Do not random erase first (clean) augmentation split')
+
+    # * Mixup params
+    parser.add_argument('--mixup', type=float, default=0.8,
+                        help='mixup alpha, mixup enabled if > 0. (default: 0.8)')
+    parser.add_argument('--cutmix', type=float, default=1.0,
+                        help='cutmix alpha, cutmix enabled if > 0. (default: 1.0)')
+    parser.add_argument('--cutmix-minmax', type=float, nargs='+', default=None,
+                        help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
+    parser.add_argument('--mixup-prob', type=float, default=1.0,
+                        help='Probability of performing mixup or cutmix when either/both is enabled')
+    parser.add_argument('--mixup-switch-prob', type=float, default=0.5,
+                        help='Probability of switching to cutmix when both mixup and cutmix enabled')
+    parser.add_argument('--mixup-mode', type=str, default='batch',
+                        help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')
+
+    # Distillation parameters
+    parser.add_argument('--teacher-model', default='regnety_160', type=str, metavar='MODEL',
+                        help='Name of teacher model to train (default: "regnety_160"')
+    parser.add_argument('--teacher-path', type=str, default='')
+    parser.add_argument('--distillation-type', default='none', choices=['none', 'soft', 'hard'], type=str, help="")
+    parser.add_argument('--distillation-alpha', default=0.5, type=float, help="")
+    parser.add_argument('--distillation-tau', default=1.0, type=float, help="")
+
+    # * Finetuning params
+    parser.add_argument('--finetune', default='', help='finetune from checkpoint')
+    parser.add_argument('--attn-only', action='store_true')
+
+    # Dataset parameters
+    parser.add_argument('--data-path', default='/datasets01/imagenet_full_size/061417/', type=str,
+                        help='dataset path')
+    parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19'],
+                        type=str, help='Image Net dataset path')
+    parser.add_argument('--inat-category', default='name',
+                        choices=['kingdom', 'phylum', 'class', 'order', 'supercategory', 'family', 'genus', 'name'],
+                        type=str, help='semantic granularity')
+
+    parser.add_argument('--output_dir', default='',
+                        help='path where to save, empty for no saving')
+    parser.add_argument('--device', default='cpu',
+                        help='device to use for training / testing')
+    parser.add_argument('--seed', default=0, type=int)
+    parser.add_argument('--resume', default='', help='resume from checkpoint')
+    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
+                        help='start epoch')
+    parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
+    parser.add_argument('--eval-crop-ratio', default=0.875, type=float, help="Crop ratio for evaluation")
+    parser.add_argument('--dist-eval', action='store_true', default=False, help='Enabling distributed evaluation')
+    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--pin-mem', action='store_true',
+                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+    parser.add_argument('--no-pin-mem', action='store_false', dest='pin_mem',
+                        help='')
+    parser.set_defaults(pin_mem=True)
+
+    # distributed training parameters
+    parser.add_argument('--world_size', default=1, type=int,
+                        help='number of distributed processes')
+    parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')
+
+    parser.add_argument('--exp_name', default='deit', type=str, help='experiment name')
+    parser.add_argument('--config', default=None, type=str, help='configuration')
+    parser.add_argument('--scoring', action='store_true', default=False, help='configuration')
+    parser.add_argument('--proxy', default='synflow', type=str, help='configuration')
+    parser.add_argument('--snnet_name', default='snnetv2', type=str, help='configuration')
+    parser.add_argument('--get_flops', action='store_true')
+    parser.add_argument('--flops_sampling_k', default=None, type=float, help="Crop ratio for evaluation")
+    parser.add_argument('--low_rank', action='store_true', default=False, help='Enabling distributed evaluation')
+    parser.add_argument('--lora_r', default=64, type=int,
+                        help='number of distributed processes')
+    parser.add_argument('--flops_gap', default=1.0, type=float,
+                        help='number of distributed processes')
+
+    return parser
+
+
+def initialize_model_stitching_layer(model, mixup_fn, data_loader, device):
+    for samples, targets in data_loader:
+        samples = samples.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+
+        if mixup_fn is not None:
+            samples, targets = mixup_fn(samples, targets)
+
+        with torch.cuda.amp.autocast():
+            model.initialize_stitching_weights(samples)
+
+        break
+
+
+@torch.no_grad()
+def analyse_flops_for_all(model, config_name):
+    all_cfgs = model.all_cfgs
+    stitch_results = {}
+
+    for cfg_id in all_cfgs:
+        model.reset_stitch_id(cfg_id)
+        flops = FlopCountAnalysis(model, torch.randn(1, 3, 224, 224).cuda()).total()
+        stitch_results[cfg_id] = flops
+
+    save_dir = './model_flops'
+    if not os.path.exists(save_dir):
+        os.mkdir(save_dir)
+
+    with open(os.path.join(save_dir, f'flops_{config_name}.json'), 'w+') as f:
+        json.dump(stitch_results, f, indent=4)
+
+
+def main(args):
+    utils.init_distributed_mode(args)
+
+    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
+    logger = get_root_logger(os.path.join(args.output_dir, f'{timestamp}.log'))
+
+    logger.info(str(args))
+
+    if args.distillation_type != 'none' and args.finetune and not args.eval:
+        raise NotImplementedError("Finetuning with distillation not yet supported")
+
+    device = torch.device(args.device)
+
+    # fix the seed for reproducibility
+    seed = args.seed + utils.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    # random.seed(seed)
+
+    cudnn.benchmark = True
+
+    from datasets import build_transform
+
+    transform = build_transform(False, args)
+
+    anchors = []
+    for i, anchor_name in enumerate(args.anchors):
+        logger.info(f"Creating model: {anchor_name}")
+        anchor = create_model(
+            anchor_name,
+            pretrained=False,
+            pretrained_deit=None,
+            num_classes=1000,
+            drop_path_rate=args.anchor_drop_path[i],
+            img_size=args.input_size
+        )
+        anchors.append(anchor)
+
+    model = SNNetv2(anchors, lora_r=args.lora_r)
+
+    checkpoint = torch.load(args.resume, map_location='cpu')
+    # torch.save({'model': checkpoint['model']}, './snnetv2_deit3_s_l_50ep.pth')
+
+    logger.info(f"load checkpoint from {args.resume}")
+    model.load_state_dict(checkpoint['model'])
+
+    model.to(device)
+
+    config_name = args.config.split('/')[-1].split('.')[0]
+    model.eval()
+
+    eval_res = {}
+    flops_res = {}
+    with open('stitches_res_s_l.txt', 'r') as f:
+        for line in f.readlines():
+            epoch_stat = json.loads(line.strip())
+            eval_res[epoch_stat['cfg_id']] = epoch_stat['acc1']
+            flops_res[epoch_stat['cfg_id']] = epoch_stat['flops'] / 1e9
+
+    def visualize_stitch_pos(stitch_id):
+        if stitch_id == 13:
+            # 13 is equivalent to 0
+            stitch_id = 0
+
+        names = [f'ID {key}' for key in flops_res.keys()]
+
+        fig = px.scatter(x=flops_res.values(), y=eval_res.values(), hover_name=names)
+        fig.update_layout(
+            title=f"SN-Netv2 - Stitch ID - {stitch_id}",
+            title_x=0.5,
+            xaxis_title="GFLOPs",
+            yaxis_title="mIoU",
+            font=dict(
+                family="Courier New, monospace",
+                size=18,
+                color="RebeccaPurple"
+            ),
+            legend=dict(
+                yanchor="bottom",
+                y=0.99,
+                xanchor="left",
+                x=0.01),
+        )
+        # continent, DarkSlateGrey
+        fig.update_traces(marker=dict(size=10,
+                                      line=dict(width=2)),
+                          selector=dict(mode='markers'))
+
+        fig.add_scatter(x=[flops_res[stitch_id]], y=[eval_res[stitch_id]], mode='markers', marker=dict(size=15),
+                        name='Current Stitch')
+        return fig
+
+    # Download human-readable labels for ImageNet.
+    response = requests.get("https://git.io/JJkYN")
+    labels = response.text.split("\n")
+
+    def process_image(image, stitch_id):
+        # inp = torch.from_numpy(image).permute(2, 0, 1).float()
+        inp = transform(image).unsqueeze(0).to(device)
+        model.reset_stitch_id(stitch_id)
+        with torch.no_grad():
+            prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
+            confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
+        fig = visualize_stitch_pos(stitch_id)
+        return confidences, fig
+
+    with gr.Blocks() as main_page:
+        with gr.Column():
+            gr.HTML("""
+                <h1 align="center" style=" display: flex; flex-direction: row; justify-content: center; font-size: 25pt; ">Stitched ViTs are Flexible Vision Backbones</h1>
+                <div align="center"> <img align="center" src='file/gradio_banner.png' width="70%"> </div>
+                <h3 align="center" >This is the classification demo page of SN-Netv2, an flexible vision backbone that allows for 100+ runtime speed and performance trade-offs.</h3>
+                <h3 align="center" >You can also run this gradio demo on your local GPUs at https://github.com/ziplab/SN-Netv2</h3>
+                """)
+            with gr.Row():
+                with gr.Column():
+                    image_input = gr.Image(type='pil')
+                    stitch_slider = gr.Slider(minimum=0, maximum=134, step=1, label="Stitch ID")
+                    with gr.Row():
+                        clear_button = gr.ClearButton()
+                        submit_button = gr.Button()
+                with gr.Column():
+                    label_output = gr.Label(num_top_classes=5)
+                    stitch_plot = gr.Plot(label='Stitch Position')
+
+        submit_button.click(
+            fn=process_image,
+            inputs=[image_input, stitch_slider],
+            outputs=[label_output, stitch_plot],
+        )
+
+        stitch_slider.change(
+            fn=visualize_stitch_pos,
+            inputs=[stitch_slider],
+            outputs=[stitch_plot],
+            show_progress=False
+        )
+
+        clear_button.click(
+            lambda: [None, 0, None, None],
+            outputs=[image_input, stitch_slider, label_output, stitch_plot],
+        )
+
+        gr.Examples(
+            [
+                ['demo.jpg', 0],
+            ],
+            inputs=[
+                image_input,
+                stitch_slider
+            ],
+            outputs=[
+                label_output,
+                stitch_plot
+            ],
+        )
+
+    main_page.launch(allowed_paths=['./'])
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('DeiT training and evaluation script', parents=[get_args_parser()])
+    args = parser.parse_args()
+    setattr(args, 'config', f'gradio_demo.json')
+    if args.config is not None:
+        config_args = json.load(open(args.config))
+        override_keys = {arg[2:].split('=')[0] for arg in sys.argv[1:]
+                         if arg.startswith('--')}
+        for k, v in config_args.items():
+            if k not in override_keys:
+                setattr(args, k, v)
+
+    output_dir = os.path.join('outputs', args.exp_name)
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    checkpoint_path = os.path.join(output_dir, 'checkpoint.pth')
+    if os.path.exists(checkpoint_path) and not args.resume:
+        setattr(args, 'resume', checkpoint_path)
 
-def greet(name):
-    return "Hello " + name + "!!"
+    setattr(args, 'output_dir', output_dir)
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+    main(args)

datasets.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+import os
+import json
+
+from torchvision import datasets, transforms
+from torchvision.datasets.folder import ImageFolder, default_loader
+
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.data import create_transform
+
+
+class INatDataset(ImageFolder):
+    def __init__(self, root, train=True, year=2018, transform=None, target_transform=None,
+                 category='name', loader=default_loader):
+        self.transform = transform
+        self.loader = loader
+        self.target_transform = target_transform
+        self.year = year
+        # assert category in ['kingdom','phylum','class','order','supercategory','family','genus','name']
+        path_json = os.path.join(root, f'{"train" if train else "val"}{year}.json')
+        with open(path_json) as json_file:
+            data = json.load(json_file)
+
+        with open(os.path.join(root, 'categories.json')) as json_file:
+            data_catg = json.load(json_file)
+
+        path_json_for_targeter = os.path.join(root, f"train{year}.json")
+
+        with open(path_json_for_targeter) as json_file:
+            data_for_targeter = json.load(json_file)
+
+        targeter = {}
+        indexer = 0
+        for elem in data_for_targeter['annotations']:
+            king = []
+            king.append(data_catg[int(elem['category_id'])][category])
+            if king[0] not in targeter.keys():
+                targeter[king[0]] = indexer
+                indexer += 1
+        self.nb_classes = len(targeter)
+
+        self.samples = []
+        for elem in data['images']:
+            cut = elem['file_name'].split('/')
+            target_current = int(cut[2])
+            path_current = os.path.join(root, cut[0], cut[2], cut[3])
+
+            categors = data_catg[target_current]
+            target_current_true = targeter[categors[category]]
+            self.samples.append((path_current, target_current_true))
+
+    # __getitem__ and __len__ inherited from ImageFolder
+
+
+def build_dataset(is_train, args):
+    transform = build_transform(is_train, args)
+
+    if args.data_set == 'CIFAR':
+        dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform)
+        nb_classes = 100
+    elif args.data_set == 'IMNET':
+        root = os.path.join(args.data_path, 'train' if is_train else 'val')
+        dataset = datasets.ImageFolder(root, transform=transform)
+        nb_classes = 1000
+    elif args.data_set == 'INAT':
+        dataset = INatDataset(args.data_path, train=is_train, year=2018,
+                              category=args.inat_category, transform=transform)
+        nb_classes = dataset.nb_classes
+    elif args.data_set == 'INAT19':
+        dataset = INatDataset(args.data_path, train=is_train, year=2019,
+                              category=args.inat_category, transform=transform)
+        nb_classes = dataset.nb_classes
+
+    return dataset, nb_classes
+
+
+def build_transform(is_train, args):
+    resize_im = args.input_size > 32
+    if is_train:
+        # this should always dispatch to transforms_imagenet_train
+        transform = create_transform(
+            input_size=args.input_size,
+            is_training=True,
+            color_jitter=args.color_jitter,
+            auto_augment=args.aa,
+            interpolation=args.train_interpolation,
+            re_prob=args.reprob,
+            re_mode=args.remode,
+            re_count=args.recount,
+        )
+        if not resize_im:
+            # replace RandomResizedCropAndInterpolation with
+            # RandomCrop
+            transform.transforms[0] = transforms.RandomCrop(
+                args.input_size, padding=4)
+        return transform
+
+    t = []
+    if resize_im:
+        size = int(args.input_size / args.eval_crop_ratio)
+        t.append(
+            transforms.Resize(size, interpolation=3),  # to maintain same ratio w.r.t. 224 images
+        )
+        t.append(transforms.CenterCrop(args.input_size))
+
+    t.append(transforms.ToTensor())
+    t.append(transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD))
+    return transforms.Compose(t)

flops_gradio_demo.json

@@ -0,0 +1,136 @@
+{
+    "0": 4608338304,
+    "1": 61604135936,
+    "2": 56843745792,
+    "3": 52102230016,
+    "4": 47360714240,
+    "5": 42619198464,
+    "6": 37877682688,
+    "7": 33136166912,
+    "8": 28394651136,
+    "9": 23653135360,
+    "10": 18911619584,
+    "11": 14170103808,
+    "12": 9428588032,
+    "14": 9523655552,
+    "15": 14265171328,
+    "16": 19006687104,
+    "17": 23748202880,
+    "18": 28489718656,
+    "19": 33231234432,
+    "20": 37972750208,
+    "21": 42714265984,
+    "22": 47455781760,
+    "23": 52197297536,
+    "24": 56938813312,
+    "25": 57017547264,
+    "26": 52276031488,
+    "27": 47534515712,
+    "28": 42792999936,
+    "29": 38051484160,
+    "30": 33309968384,
+    "31": 28568452608,
+    "32": 23826936832,
+    "33": 19085421056,
+    "34": 14343905280,
+    "35": 57017547264,
+    "36": 52276031488,
+    "37": 47534515712,
+    "38": 42792999936,
+    "39": 38051484160,
+    "40": 33309968384,
+    "41": 28568452608,
+    "42": 23826936832,
+    "43": 19085421056,
+    "44": 57017547264,
+    "45": 52276031488,
+    "46": 47534515712,
+    "47": 42792999936,
+    "48": 38051484160,
+    "49": 33309968384,
+    "50": 28568452608,
+    "51": 23826936832,
+    "52": 57017547264,
+    "53": 52276031488,
+    "54": 47534515712,
+    "55": 42792999936,
+    "56": 38051484160,
+    "57": 33309968384,
+    "58": 28568452608,
+    "59": 57017547264,
+    "60": 52276031488,
+    "61": 47534515712,
+    "62": 42792999936,
+    "63": 38051484160,
+    "64": 33309968384,
+    "65": 57017547264,
+    "66": 52276031488,
+    "67": 47534515712,
+    "68": 42792999936,
+    "69": 38051484160,
+    "70": 57017547264,
+    "71": 52276031488,
+    "72": 47534515712,
+    "73": 42792999936,
+    "74": 57017547264,
+    "75": 52276031488,
+    "76": 47534515712,
+    "77": 57017547264,
+    "78": 52276031488,
+    "79": 57017547264,
+    "80": 9504781184,
+    "81": 14246296960,
+    "82": 18987812736,
+    "83": 23729328512,
+    "84": 28470844288,
+    "85": 33212360064,
+    "86": 37953875840,
+    "87": 42695391616,
+    "88": 47436907392,
+    "89": 52178423168,
+    "90": 9504781184,
+    "91": 14246296960,
+    "92": 18987812736,
+    "93": 23729328512,
+    "94": 28470844288,
+    "95": 33212360064,
+    "96": 37953875840,
+    "97": 42695391616,
+    "98": 47436907392,
+    "99": 9504781184,
+    "100": 14246296960,
+    "101": 18987812736,
+    "102": 23729328512,
+    "103": 28470844288,
+    "104": 33212360064,
+    "105": 37953875840,
+    "106": 42695391616,
+    "107": 9504781184,
+    "108": 14246296960,
+    "109": 18987812736,
+    "110": 23729328512,
+    "111": 28470844288,
+    "112": 33212360064,
+    "113": 37953875840,
+    "114": 9504781184,
+    "115": 14246296960,
+    "116": 18987812736,
+    "117": 23729328512,
+    "118": 28470844288,
+    "119": 33212360064,
+    "120": 9504781184,
+    "121": 14246296960,
+    "122": 18987812736,
+    "123": 23729328512,
+    "124": 28470844288,
+    "125": 9504781184,
+    "126": 14246296960,
+    "127": 18987812736,
+    "128": 23729328512,
+    "129": 9504781184,
+    "130": 14246296960,
+    "131": 18987812736,
+    "132": 9504781184,
+    "133": 14246296960,
+    "134": 9504781184
+}

gradio_demo.json

@@ -0,0 +1,33 @@
+{
+  "anchors": ["deit_small_patch16_LS", "deit_large_patch16_LS"],
+  "batch_size": 64,
+  "snnet_name": "snnet_v2",
+  "data_path": "/data2/datasets/imagenet",
+  "data_set": "IMNET",
+  "exp_name": "stitch_s_l_v2_lora_r_64_50_ep",
+  "input_size": 224,
+  "num_workers": 10,
+  "lr": 0.00003,
+  "warmup_lr": 1e-7,
+  "epochs": 50,
+  "weight_decay": 0.02,
+  "sched": "cosine",
+  "eval_crop_ratio": 1.0,
+  "reprob": 0.0,
+  "smoothing": 0.1,
+  "warmup_epochs": 5,
+  "drop": 0.0,
+  "seed": 0,
+  "opt": "fusedlamb",
+  "mixup": 0,
+  "anchor_drop_path": [0.05, 0.4],
+  "cutmix": 1.0,
+  "color_jitter": 0.3,
+  "unscale_lr": true,
+  "no_repeated_aug": true,
+  "ThreeAugment": true,
+  "src": true,
+  "lora_r": 64,
+  "pretrained_deit": "../pretrained_weights",
+  "resume": "snnetv2_deit3_s_l.pth"
+}

models_v2.py

@@ -0,0 +1,568 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+import os.path
+
+import torch
+import torch.nn as nn
+from functools import partial
+
+from timm.models.vision_transformer import Mlp, PatchEmbed , _cfg
+
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models.registry import register_model
+# from xformers.ops import memory_efficient_attention
+
+class Attention(nn.Module):
+    # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+
+        # x = memory_efficient_attention(q, k, v).transpose(1, 2).reshape(B, N, C)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+    
+class Block(nn.Module):
+    # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,Attention_block = Attention,Mlp_block=Mlp
+                 ,init_values=1e-4):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention_block(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x 
+    
+class Layer_scale_init_Block(nn.Module):
+    # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+    # with slight modifications
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,Attention_block = Attention,Mlp_block=Mlp
+                 ,init_values=1e-4):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention_block(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+class Layer_scale_init_Block_paralx2(nn.Module):
+    # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+    # with slight modifications
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,Attention_block = Attention,Mlp_block=Mlp
+                 ,init_values=1e-4):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.norm11 = norm_layer(dim)
+        self.attn = Attention_block(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.attn1 = Attention_block(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        self.norm21 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.mlp1 = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        self.gamma_1_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        self.gamma_2_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        
+    def forward(self, x):
+        x = x + self.drop_path(self.gamma_1*self.attn(self.norm1(x))) + self.drop_path(self.gamma_1_1 * self.attn1(self.norm11(x)))
+        x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x))) + self.drop_path(self.gamma_2_1 * self.mlp1(self.norm21(x)))
+        return x
+        
+class Block_paralx2(nn.Module):
+    # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+    # with slight modifications
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,Attention_block = Attention,Mlp_block=Mlp
+                 ,init_values=1e-4):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.norm11 = norm_layer(dim)
+        self.attn = Attention_block(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.attn1 = Attention_block(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        self.norm21 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.mlp1 = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x))) + self.drop_path(self.attn1(self.norm11(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x))) + self.drop_path(self.mlp1(self.norm21(x)))
+        return x
+        
+        
+class hMLP_stem(nn.Module):
+    """ hMLP_stem: https://arxiv.org/pdf/2203.09795.pdf
+    taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+    with slight modifications
+    """
+    def __init__(self, img_size=224,  patch_size=16, in_chans=3, embed_dim=768,norm_layer=nn.SyncBatchNorm):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+        self.proj = torch.nn.Sequential(*[nn.Conv2d(in_chans, embed_dim//4, kernel_size=4, stride=4),
+                                          norm_layer(embed_dim//4),
+                                          nn.GELU(),
+                                          nn.Conv2d(embed_dim//4, embed_dim//4, kernel_size=2, stride=2),
+                                          norm_layer(embed_dim//4),
+                                          nn.GELU(),
+                                          nn.Conv2d(embed_dim//4, embed_dim, kernel_size=2, stride=2),
+                                          norm_layer(embed_dim),
+                                         ])
+        
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+    
+class vit_models(nn.Module):
+    """ Vision Transformer with LayerScale (https://arxiv.org/abs/2103.17239) support
+    taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+    with slight modifications
+    """
+    def __init__(self, img_size=224,  patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, global_pool=None,
+                 block_layers = Block,
+                 Patch_layer=PatchEmbed,act_layer=nn.GELU,
+                 Attention_block = Attention, Mlp_block=Mlp,
+                dpr_constant=True,init_scale=1e-4,
+                mlp_ratio_clstk = 4.0):
+        super().__init__()
+        
+        self.dropout_rate = drop_rate
+        self.depth = depth
+
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim
+
+        self.patch_embed = Patch_layer(
+                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+
+        dpr = [drop_path_rate for i in range(depth)]
+        self.blocks = nn.ModuleList([
+            block_layers(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=0.0, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                act_layer=act_layer,Attention_block=Attention_block,Mlp_block=Mlp_block,init_values=init_scale)
+            for i in range(depth)])
+
+        self.norm = norm_layer(embed_dim)
+
+        self.feature_info = [dict(num_chs=embed_dim, reduction=0, module='head')]
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def get_classifier(self):
+        return self.head
+    
+    def get_num_layers(self):
+        return len(self.blocks)
+    
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+
+    def extract_block_features(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+
+        x = x + self.pos_embed
+
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        outs = {}
+
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            outs[i] = x.detach()
+        return outs
+
+    def selective_forward(self, x, begin, end):
+        for i, blk in enumerate(self.blocks):
+            if i < begin:
+                continue
+            if i > end:
+                break
+            x = blk(x)
+        return x
+
+    def forward_until(self, x, blk_id):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = x + self.pos_embed
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i == blk_id:
+                break
+
+        return x
+
+    def forward_from(self, x, blk_id):
+        for i, blk in enumerate(self.blocks):
+            if i < blk_id:
+                continue
+            x = blk(x)
+
+        x = self.norm(x)
+        x = self.head(x[:, 0])
+
+        return x
+
+    def forward_patch_embed(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        
+        x = x + self.pos_embed
+        
+        x = torch.cat((cls_tokens, x), dim=1)
+        return x
+
+    
+    def forward_norm_head(self, x):
+        x = self.norm(x)
+        x = self.head(x[:, 0])
+        return x
+
+
+    def forward_features(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        
+        x = x + self.pos_embed
+        
+        x = torch.cat((cls_tokens, x), dim=1)
+            
+        for i , blk in enumerate(self.blocks):
+            x = blk(x)
+            
+        x = self.norm(x)
+        return x[:, 0]
+
+    def forward(self, x):
+
+        x = self.forward_features(x)
+        
+        if self.dropout_rate:
+            x = F.dropout(x, p=float(self.dropout_rate), training=self.training)
+        x = self.head(x)
+        
+        return x
+
+# DeiT III: Revenge of the ViT (https://arxiv.org/abs/2204.07118)
+
+@register_model
+def deit_tiny_patch16_LS(pretrained=False, img_size=224, pretrained_21k = False,  pretrained_cfg_overlay=None, **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Layer_scale_init_Block, **kwargs)
+    
+    return model
+    
+    
+@register_model
+def deit_small_patch16_LS(pretrained=False, img_size=224, pretrained_21k = False, pretrained_cfg=None, pretrained_deit=None, pretrained_cfg_overlay=None, **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Layer_scale_init_Block, **kwargs)
+    model.default_cfg = _cfg()
+    if pretrained:
+        # name = 'https://dl.fbaipublicfiles.com/deit/deit_3_small_'+str(img_size)+'_'
+        # if pretrained_21k:
+        #     name+='21k.pth'
+        # else:
+        #     name+='1k.pth'
+            
+        # checkpoint = torch.hub.load_state_dict_from_url(
+        #     url=name,
+        #     map_location="cpu", check_hash=True
+        # )
+        checkpoint = torch.load(os.path.join(pretrained_deit, 'deit_3_small_224_21k.pth'))
+        model.load_state_dict(checkpoint["model"])
+
+    return model
+
+@register_model
+def deit_medium_patch16_LS(pretrained=False, img_size=224, pretrained_21k = False, **kwargs):
+    model = vit_models(
+        patch_size=16, embed_dim=512, depth=12, num_heads=8, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers = Layer_scale_init_Block, **kwargs)
+    model.default_cfg = _cfg()
+    if pretrained:
+        name = 'https://dl.fbaipublicfiles.com/deit/deit_3_medium_'+str(img_size)+'_'
+        if pretrained_21k:
+            name+='21k.pth'
+        else:
+            name+='1k.pth'
+            
+        checkpoint = torch.hub.load_state_dict_from_url(
+            url=name,
+            map_location="cpu", check_hash=True
+        )
+        model.load_state_dict(checkpoint["model"])
+    return model 
+
+@register_model
+def deit_base_patch16_LS(pretrained=False, pretrained_cfg=None, img_size=224, pretrained_21k = False,  pretrained_deit=None, pretrained_cfg_overlay=None, **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Layer_scale_init_Block, **kwargs)
+    if pretrained:
+        # name = 'https://dl.fbaipublicfiles.com/deit/deit_3_small_'+str(img_size)+'_'
+        # if pretrained_21k:
+        #     name+='21k.pth'
+        # else:
+        #     name+='1k.pth'
+
+        # checkpoint = torch.hub.load_state_dict_from_url(
+        #     url=name,
+        #     map_location="cpu", check_hash=True
+        # )
+        checkpoint = torch.load(os.path.join(pretrained_deit, 'deit_3_base_224_21k.pth'))
+        model.load_state_dict(checkpoint["model"])
+    return model
+    
+@register_model
+def deit_large_patch16_LS(pretrained=False, img_size=224, pretrained_21k = False,  pretrained_cfg=None, pretrained_deit=None, pretrained_cfg_overlay=None, **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Layer_scale_init_Block, **kwargs)
+    if pretrained:
+        # name = 'https://dl.fbaipublicfiles.com/deit/deit_3_large_'+str(img_size)+'_'
+        # if pretrained_21k:
+        #     name+='21k.pth'
+        # else:
+        #     name+='1k.pth'
+        #
+        # checkpoint = torch.hub.load_state_dict_from_url(
+        #     url=name,
+        #     map_location="cpu", check_hash=True
+        # )
+        checkpoint = torch.load(os.path.join(pretrained_deit, 'deit_3_large_224_21k.pth'))
+        model.load_state_dict(checkpoint["model"])
+    return model
+    
+@register_model
+def deit_huge_patch14_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=14, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers = Layer_scale_init_Block, **kwargs)
+    if pretrained:
+        name = 'https://dl.fbaipublicfiles.com/deit/deit_3_huge_'+str(img_size)+'_'
+        if pretrained_21k:
+            name+='21k_v1.pth'
+        else:
+            name+='1k_v1.pth'
+            
+        checkpoint = torch.hub.load_state_dict_from_url(
+            url=name,
+            map_location="cpu", check_hash=True
+        )
+        model.load_state_dict(checkpoint["model"])
+    return model
+    
+@register_model
+def deit_huge_patch14_52_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=14, embed_dim=1280, depth=52, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers = Layer_scale_init_Block, **kwargs)
+
+    return model
+    
+@register_model
+def deit_huge_patch14_26x2_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=14, embed_dim=1280, depth=26, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers = Layer_scale_init_Block_paralx2, **kwargs)
+
+    return model
+    
+@register_model
+def deit_Giant_48x2_patch14_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=14, embed_dim=1664, depth=48, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers = Block_paral_LS, **kwargs)
+
+    return model
+
+@register_model
+def deit_giant_40x2_patch14_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=14, embed_dim=1408, depth=40, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers = Block_paral_LS, **kwargs)
+    return model
+
+@register_model
+def deit_Giant_48_patch14_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=14, embed_dim=1664, depth=48, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers = Layer_scale_init_Block, **kwargs)
+    return model
+
+@register_model
+def deit_giant_40_patch14_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=14, embed_dim=1408, depth=40, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers = Layer_scale_init_Block, **kwargs)
+    #model.default_cfg = _cfg()
+
+    return model
+
+# Models from Three things everyone should know about Vision Transformers (https://arxiv.org/pdf/2203.09795.pdf)
+
+@register_model
+def deit_small_patch16_36_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=384, depth=36, num_heads=6, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Layer_scale_init_Block, **kwargs)
+
+    return model
+    
+@register_model
+def deit_small_patch16_36(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=384, depth=36, num_heads=6, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+
+    return model
+    
+@register_model
+def deit_small_patch16_18x2_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=384, depth=18, num_heads=6, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Layer_scale_init_Block_paralx2, **kwargs)
+
+    return model
+    
+@register_model
+def deit_small_patch16_18x2(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=384, depth=18, num_heads=6, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Block_paralx2, **kwargs)
+
+    return model
+    
+  
+@register_model
+def deit_base_patch16_18x2_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=768, depth=18, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Layer_scale_init_Block_paralx2, **kwargs)
+
+    return model
+
+
+@register_model
+def deit_base_patch16_18x2(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=768, depth=18, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Block_paralx2, **kwargs)
+
+    return model
+    
+
+@register_model
+def deit_base_patch16_36x1_LS(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=768, depth=36, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),block_layers=Layer_scale_init_Block, **kwargs)
+
+    return model
+
+@register_model
+def deit_base_patch16_36x1(pretrained=False, img_size=224, pretrained_21k = False,  **kwargs):
+    model = vit_models(
+        img_size = img_size, patch_size=16, embed_dim=768, depth=36, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+
+    return model
+

outputs/deit/20240118_171921.log

@@ -0,0 +1 @@
+2024-01-18 17:19:21,866 - snnet - INFO - Namespace(batch_size=64, epochs=300, bce_loss=False, unscale_lr=False, model='deit_base_patch16_224', input_size=224, drop=0.0, drop_path=0.1, model_ema=True, model_ema_decay=0.99996, model_ema_force_cpu=False, opt='adamw', opt_eps=1e-08, opt_betas=None, clip_grad=None, momentum=0.9, weight_decay=0.05, sched='cosine', lr=0.0005, lr_noise=None, lr_noise_pct=0.67, lr_noise_std=1.0, warmup_lr=1e-06, min_lr=1e-05, decay_epochs=30, warmup_epochs=5, cooldown_epochs=10, patience_epochs=10, decay_rate=0.1, color_jitter=0.3, aa='rand-m9-mstd0.5-inc1', smoothing=0.1, train_interpolation='bicubic', repeated_aug=True, train_mode=True, ThreeAugment=False, src=False, reprob=0.25, remode='pixel', recount=1, resplit=False, mixup=0.8, cutmix=1.0, cutmix_minmax=None, mixup_prob=1.0, mixup_switch_prob=0.5, mixup_mode='batch', teacher_model='regnety_160', teacher_path='', distillation_type='none', distillation_alpha=0.5, distillation_tau=1.0, finetune='', attn_only=False, data_path='/datasets01/imagenet_full_size/061417/', data_set='IMNET', inat_category='name', output_dir='outputs/deit', device='cuda', seed=0, resume='', start_epoch=0, eval=False, eval_crop_ratio=0.875, dist_eval=False, num_workers=10, pin_mem=True, world_size=1, dist_url='env://', exp_name='deit', config=None, scoring=False, proxy='synflow', snnet_name='snnetv2', get_flops=False, flops_sampling_k=None, low_rank=False, lora_r=64, flops_gap=1.0, distributed=False)

outputs/stitch_s_l_v2_lora_r_64_50_ep/20240118_172124.log

@@ -0,0 +1,2 @@
+2024-01-18 17:21:24,162 - snnet - INFO - Namespace(batch_size=64, epochs=50, bce_loss=False, unscale_lr=True, model='deit_base_patch16_224', input_size=224, drop=0.0, drop_path=0.1, model_ema=True, model_ema_decay=0.99996, model_ema_force_cpu=False, opt='fusedlamb', opt_eps=1e-08, opt_betas=None, clip_grad=None, momentum=0.9, weight_decay=0.02, sched='cosine', lr=3e-05, lr_noise=None, lr_noise_pct=0.67, lr_noise_std=1.0, warmup_lr=1e-07, min_lr=1e-05, decay_epochs=30, warmup_epochs=5, cooldown_epochs=10, patience_epochs=10, decay_rate=0.1, color_jitter=0.3, aa='rand-m9-mstd0.5-inc1', smoothing=0.1, train_interpolation='bicubic', repeated_aug=True, train_mode=True, ThreeAugment=True, src=True, reprob=0.0, remode='pixel', recount=1, resplit=False, mixup=0, cutmix=1.0, cutmix_minmax=None, mixup_prob=1.0, mixup_switch_prob=0.5, mixup_mode='batch', teacher_model='regnety_160', teacher_path='', distillation_type='none', distillation_alpha=0.5, distillation_tau=1.0, finetune='', attn_only=False, data_path='/data2/datasets/imagenet', data_set='IMNET', inat_category='name', output_dir='outputs/stitch_s_l_v2_lora_r_64_50_ep', device='cuda', seed=0, resume='snnetv2_deit3_s_l.pth', start_epoch=0, eval=False, eval_crop_ratio=1.0, dist_eval=False, num_workers=10, pin_mem=True, world_size=1, dist_url='env://', exp_name='stitch_s_l_v2_lora_r_64_50_ep', config='gradio_demo.json', scoring=False, proxy='synflow', snnet_name='snnet_v2', get_flops=False, flops_sampling_k=None, low_rank=False, lora_r=64, flops_gap=1.0, anchors=['deit_small_patch16_LS', 'deit_large_patch16_LS'], anchor_drop_path=[0.05, 0.4], no_repeated_aug=True, pretrained_deit='../pretrained_weights', distributed=False)
+2024-01-18 17:21:24,163 - snnet - INFO - Creating model: deit_small_patch16_LS

outputs/stitch_s_l_v2_lora_r_64_50_ep/20240118_172140.log

@@ -0,0 +1,2 @@
+2024-01-18 17:21:40,831 - snnet - INFO - Namespace(batch_size=64, epochs=50, bce_loss=False, unscale_lr=True, model='deit_base_patch16_224', input_size=224, drop=0.0, drop_path=0.1, model_ema=True, model_ema_decay=0.99996, model_ema_force_cpu=False, opt='fusedlamb', opt_eps=1e-08, opt_betas=None, clip_grad=None, momentum=0.9, weight_decay=0.02, sched='cosine', lr=3e-05, lr_noise=None, lr_noise_pct=0.67, lr_noise_std=1.0, warmup_lr=1e-07, min_lr=1e-05, decay_epochs=30, warmup_epochs=5, cooldown_epochs=10, patience_epochs=10, decay_rate=0.1, color_jitter=0.3, aa='rand-m9-mstd0.5-inc1', smoothing=0.1, train_interpolation='bicubic', repeated_aug=True, train_mode=True, ThreeAugment=True, src=True, reprob=0.0, remode='pixel', recount=1, resplit=False, mixup=0, cutmix=1.0, cutmix_minmax=None, mixup_prob=1.0, mixup_switch_prob=0.5, mixup_mode='batch', teacher_model='regnety_160', teacher_path='', distillation_type='none', distillation_alpha=0.5, distillation_tau=1.0, finetune='', attn_only=False, data_path='/data2/datasets/imagenet', data_set='IMNET', inat_category='name', output_dir='outputs/stitch_s_l_v2_lora_r_64_50_ep', device='cuda', seed=0, resume='snnetv2_deit3_s_l.pth', start_epoch=0, eval=False, eval_crop_ratio=1.0, dist_eval=False, num_workers=10, pin_mem=True, world_size=1, dist_url='env://', exp_name='stitch_s_l_v2_lora_r_64_50_ep', config='gradio_demo.json', scoring=False, proxy='synflow', snnet_name='snnet_v2', get_flops=False, flops_sampling_k=None, low_rank=False, lora_r=64, flops_gap=1.0, anchors=['deit_small_patch16_LS', 'deit_large_patch16_LS'], anchor_drop_path=[0.05, 0.4], no_repeated_aug=True, pretrained_deit='../pretrained_weights', distributed=False)
+2024-01-18 17:21:40,832 - snnet - INFO - Creating model: deit_small_patch16_LS

outputs/stitch_s_l_v2_lora_r_64_50_ep/20240118_172156.log

@@ -0,0 +1,5 @@
+2024-01-18 17:21:56,859 - snnet - INFO - Namespace(batch_size=64, epochs=50, bce_loss=False, unscale_lr=True, model='deit_base_patch16_224', input_size=224, drop=0.0, drop_path=0.1, model_ema=True, model_ema_decay=0.99996, model_ema_force_cpu=False, opt='fusedlamb', opt_eps=1e-08, opt_betas=None, clip_grad=None, momentum=0.9, weight_decay=0.02, sched='cosine', lr=3e-05, lr_noise=None, lr_noise_pct=0.67, lr_noise_std=1.0, warmup_lr=1e-07, min_lr=1e-05, decay_epochs=30, warmup_epochs=5, cooldown_epochs=10, patience_epochs=10, decay_rate=0.1, color_jitter=0.3, aa='rand-m9-mstd0.5-inc1', smoothing=0.1, train_interpolation='bicubic', repeated_aug=True, train_mode=True, ThreeAugment=True, src=True, reprob=0.0, remode='pixel', recount=1, resplit=False, mixup=0, cutmix=1.0, cutmix_minmax=None, mixup_prob=1.0, mixup_switch_prob=0.5, mixup_mode='batch', teacher_model='regnety_160', teacher_path='', distillation_type='none', distillation_alpha=0.5, distillation_tau=1.0, finetune='', attn_only=False, data_path='/data2/datasets/imagenet', data_set='IMNET', inat_category='name', output_dir='outputs/stitch_s_l_v2_lora_r_64_50_ep', device='cuda', seed=0, resume='snnetv2_deit3_s_l.pth', start_epoch=0, eval=False, eval_crop_ratio=1.0, dist_eval=False, num_workers=10, pin_mem=True, world_size=1, dist_url='env://', exp_name='stitch_s_l_v2_lora_r_64_50_ep', config='gradio_demo.json', scoring=False, proxy='synflow', snnet_name='snnet_v2', get_flops=False, flops_sampling_k=None, low_rank=False, lora_r=64, flops_gap=1.0, anchors=['deit_small_patch16_LS', 'deit_large_patch16_LS'], anchor_drop_path=[0.05, 0.4], no_repeated_aug=True, pretrained_deit='../pretrained_weights', distributed=False)
+2024-01-18 17:21:56,859 - snnet - INFO - Creating model: deit_small_patch16_LS
+2024-01-18 17:21:57,078 - snnet - INFO - Creating model: deit_large_patch16_LS
+2024-01-18 17:21:59,994 - snnet - INFO - [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134]
+2024-01-18 17:22:00,521 - snnet - INFO - load checkpoint from snnetv2_deit3_s_l.pth

outputs/stitch_s_l_v2_lora_r_64_50_ep/20240118_172250.log

@@ -0,0 +1,5 @@
+2024-01-18 17:22:50,304 - snnet - INFO - Namespace(batch_size=64, epochs=50, bce_loss=False, unscale_lr=True, model='deit_base_patch16_224', input_size=224, drop=0.0, drop_path=0.1, model_ema=True, model_ema_decay=0.99996, model_ema_force_cpu=False, opt='fusedlamb', opt_eps=1e-08, opt_betas=None, clip_grad=None, momentum=0.9, weight_decay=0.02, sched='cosine', lr=3e-05, lr_noise=None, lr_noise_pct=0.67, lr_noise_std=1.0, warmup_lr=1e-07, min_lr=1e-05, decay_epochs=30, warmup_epochs=5, cooldown_epochs=10, patience_epochs=10, decay_rate=0.1, color_jitter=0.3, aa='rand-m9-mstd0.5-inc1', smoothing=0.1, train_interpolation='bicubic', repeated_aug=True, train_mode=True, ThreeAugment=True, src=True, reprob=0.0, remode='pixel', recount=1, resplit=False, mixup=0, cutmix=1.0, cutmix_minmax=None, mixup_prob=1.0, mixup_switch_prob=0.5, mixup_mode='batch', teacher_model='regnety_160', teacher_path='', distillation_type='none', distillation_alpha=0.5, distillation_tau=1.0, finetune='', attn_only=False, data_path='/data2/datasets/imagenet', data_set='IMNET', inat_category='name', output_dir='outputs/stitch_s_l_v2_lora_r_64_50_ep', device='cpu', seed=0, resume='snnetv2_deit3_s_l.pth', start_epoch=0, eval=False, eval_crop_ratio=1.0, dist_eval=False, num_workers=10, pin_mem=True, world_size=1, dist_url='env://', exp_name='stitch_s_l_v2_lora_r_64_50_ep', config='gradio_demo.json', scoring=False, proxy='synflow', snnet_name='snnet_v2', get_flops=False, flops_sampling_k=None, low_rank=False, lora_r=64, flops_gap=1.0, anchors=['deit_small_patch16_LS', 'deit_large_patch16_LS'], anchor_drop_path=[0.05, 0.4], no_repeated_aug=True, pretrained_deit='../pretrained_weights', distributed=False)
+2024-01-18 17:22:50,305 - snnet - INFO - Creating model: deit_small_patch16_LS
+2024-01-18 17:22:50,535 - snnet - INFO - Creating model: deit_large_patch16_LS
+2024-01-18 17:22:53,873 - snnet - INFO - [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134]
+2024-01-18 17:22:54,392 - snnet - INFO - load checkpoint from snnetv2_deit3_s_l.pth

outputs/stitch_s_l_v2_lora_r_64_50_ep/20240118_172309.log

@@ -0,0 +1,5 @@
+2024-01-18 17:23:09,551 - snnet - INFO - Namespace(batch_size=64, epochs=50, bce_loss=False, unscale_lr=True, model='deit_base_patch16_224', input_size=224, drop=0.0, drop_path=0.1, model_ema=True, model_ema_decay=0.99996, model_ema_force_cpu=False, opt='fusedlamb', opt_eps=1e-08, opt_betas=None, clip_grad=None, momentum=0.9, weight_decay=0.02, sched='cosine', lr=3e-05, lr_noise=None, lr_noise_pct=0.67, lr_noise_std=1.0, warmup_lr=1e-07, min_lr=1e-05, decay_epochs=30, warmup_epochs=5, cooldown_epochs=10, patience_epochs=10, decay_rate=0.1, color_jitter=0.3, aa='rand-m9-mstd0.5-inc1', smoothing=0.1, train_interpolation='bicubic', repeated_aug=True, train_mode=True, ThreeAugment=True, src=True, reprob=0.0, remode='pixel', recount=1, resplit=False, mixup=0, cutmix=1.0, cutmix_minmax=None, mixup_prob=1.0, mixup_switch_prob=0.5, mixup_mode='batch', teacher_model='regnety_160', teacher_path='', distillation_type='none', distillation_alpha=0.5, distillation_tau=1.0, finetune='', attn_only=False, data_path='/data2/datasets/imagenet', data_set='IMNET', inat_category='name', output_dir='outputs/stitch_s_l_v2_lora_r_64_50_ep', device='cpu', seed=0, resume='snnetv2_deit3_s_l.pth', start_epoch=0, eval=False, eval_crop_ratio=1.0, dist_eval=False, num_workers=10, pin_mem=True, world_size=1, dist_url='env://', exp_name='stitch_s_l_v2_lora_r_64_50_ep', config='gradio_demo.json', scoring=False, proxy='synflow', snnet_name='snnet_v2', get_flops=False, flops_sampling_k=None, low_rank=False, lora_r=64, flops_gap=1.0, anchors=['deit_small_patch16_LS', 'deit_large_patch16_LS'], anchor_drop_path=[0.05, 0.4], no_repeated_aug=True, pretrained_deit='../pretrained_weights', distributed=False)
+2024-01-18 17:23:09,553 - snnet - INFO - Creating model: deit_small_patch16_LS
+2024-01-18 17:23:09,778 - snnet - INFO - Creating model: deit_large_patch16_LS
+2024-01-18 17:23:13,077 - snnet - INFO - [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134]
+2024-01-18 17:23:13,587 - snnet - INFO - load checkpoint from snnetv2_deit3_s_l.pth

outputs/stitch_s_l_v2_lora_r_64_50_ep/20240118_172332.log

@@ -0,0 +1,5 @@
+2024-01-18 17:23:32,357 - snnet - INFO - Namespace(batch_size=64, epochs=50, bce_loss=False, unscale_lr=True, model='deit_base_patch16_224', input_size=224, drop=0.0, drop_path=0.1, model_ema=True, model_ema_decay=0.99996, model_ema_force_cpu=False, opt='fusedlamb', opt_eps=1e-08, opt_betas=None, clip_grad=None, momentum=0.9, weight_decay=0.02, sched='cosine', lr=3e-05, lr_noise=None, lr_noise_pct=0.67, lr_noise_std=1.0, warmup_lr=1e-07, min_lr=1e-05, decay_epochs=30, warmup_epochs=5, cooldown_epochs=10, patience_epochs=10, decay_rate=0.1, color_jitter=0.3, aa='rand-m9-mstd0.5-inc1', smoothing=0.1, train_interpolation='bicubic', repeated_aug=True, train_mode=True, ThreeAugment=True, src=True, reprob=0.0, remode='pixel', recount=1, resplit=False, mixup=0, cutmix=1.0, cutmix_minmax=None, mixup_prob=1.0, mixup_switch_prob=0.5, mixup_mode='batch', teacher_model='regnety_160', teacher_path='', distillation_type='none', distillation_alpha=0.5, distillation_tau=1.0, finetune='', attn_only=False, data_path='/data2/datasets/imagenet', data_set='IMNET', inat_category='name', output_dir='outputs/stitch_s_l_v2_lora_r_64_50_ep', device='cpu', seed=0, resume='snnetv2_deit3_s_l.pth', start_epoch=0, eval=False, eval_crop_ratio=1.0, dist_eval=False, num_workers=10, pin_mem=True, world_size=1, dist_url='env://', exp_name='stitch_s_l_v2_lora_r_64_50_ep', config='gradio_demo.json', scoring=False, proxy='synflow', snnet_name='snnet_v2', get_flops=False, flops_sampling_k=None, low_rank=False, lora_r=64, flops_gap=1.0, anchors=['deit_small_patch16_LS', 'deit_large_patch16_LS'], anchor_drop_path=[0.05, 0.4], no_repeated_aug=True, pretrained_deit='../pretrained_weights', distributed=False)
+2024-01-18 17:23:32,358 - snnet - INFO - Creating model: deit_small_patch16_LS
+2024-01-18 17:23:32,606 - snnet - INFO - Creating model: deit_large_patch16_LS
+2024-01-18 17:23:35,576 - snnet - INFO - [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134]
+2024-01-18 17:23:36,120 - snnet - INFO - load checkpoint from snnetv2_deit3_s_l.pth

requirements.txt

@@ -0,0 +1,3 @@
+torch
+timm==0.6.12
+fvcore

snnet.py

@@ -0,0 +1,473 @@
+# Copyright (c) OpenMMLab. All rights reserved.import math
+import json
+import math
+
+import torch
+import torch.nn as nn
+
+import numpy as np
+
+from collections import defaultdict
+from utils import get_root_logger
+import torch.nn.functional as F
+
+def rearrange_activations(activations):
+    n_channels = activations.shape[-1]
+    activations = activations.reshape(-1, n_channels)
+    return activations
+
+
+def ps_inv(x1, x2):
+    '''Least-squares solver given feature maps from two anchors.
+    '''
+    x1 = rearrange_activations(x1)
+    x2 = rearrange_activations(x2)
+
+    if not x1.shape[0] == x2.shape[0]:
+        raise ValueError('Spatial size of compared neurons must match when ' \
+                         'calculating psuedo inverse matrix.')
+
+    # Get transformation matrix shape
+    shape = list(x1.shape)
+    shape[-1] += 1
+
+    # Calculate pseudo inverse
+    x1_ones = torch.ones(shape)
+    x1_ones[:, :-1] = x1
+    A_ones = torch.matmul(torch.linalg.pinv(x1_ones), x2.to(x1_ones.device)).T
+
+    # Get weights and bias
+    w = A_ones[..., :-1]
+    b = A_ones[..., -1]
+
+    return w, b
+
+
+def reset_out_indices(front_depth=12, end_depth=24, out_indices=(9, 14, 19, 23)):
+    block_ids = torch.tensor(list(range(front_depth)))
+    block_ids = block_ids[None, None, :].float()
+    end_mapping_ids = torch.nn.functional.interpolate(block_ids, end_depth)
+    end_mapping_ids = end_mapping_ids.squeeze().long().tolist()
+
+    small_out_indices = []
+    for i, idx in enumerate(end_mapping_ids):
+        if i in out_indices:
+            small_out_indices.append(idx)
+
+    return small_out_indices
+
+
+def get_stitch_configs_general_unequal(depths):
+    depths = sorted(depths)
+
+    total_configs = []
+
+    # anchor configurations
+    total_configs.append({'comb_id': [1], })
+    num_stitches = depths[0]
+    for i, blk_id in enumerate(range(num_stitches)):
+        total_configs.append({
+            'comb_id': (0, 1),
+            'stitch_cfgs': (i, (i + 1) * (depths[1] // depths[0]))
+        })
+    return total_configs, num_stitches
+
+def get_stitch_configs_bidirection(depths):
+    depths = sorted(depths)
+
+    total_configs = []
+
+    # anchor configurations
+    total_configs.append({'comb_id': [0], })
+    total_configs.append({'comb_id': [1], })
+
+    num_stitches = depths[0]
+
+    # small --> large
+    sl_configs = []
+    for i, blk_id in enumerate(range(num_stitches)):
+        sl_configs.append({
+            'comb_id': [0, 1],
+            'stitch_cfgs': [
+                [i, (i + 1) * (depths[1] // depths[0])]
+            ],
+            'stitch_layer_ids': [i]
+        })
+
+    ls_configs = []
+    lsl_confgs = []
+    block_ids = torch.tensor(list(range(depths[0])))
+    block_ids = block_ids[None, None, :].float()
+    end_mapping_ids = torch.nn.functional.interpolate(block_ids, depths[1])
+    end_mapping_ids = end_mapping_ids.squeeze().long().tolist()
+
+    # large --> small
+    for i in range(depths[1]):
+        if depths[1] != depths[0]:
+            if i % 2 == 1 and i < (depths[1] - 1):
+                ls_configs.append({
+                    'comb_id': [1, 0],
+                    'stitch_cfgs': [[i, end_mapping_ids[i] + 1]],
+                    'stitch_layer_ids': [i // (depths[1] // depths[0])]
+                })
+        else:
+            if i < (depths[1] - 1):
+                ls_configs.append({
+                    'comb_id': [1, 0],
+                    'stitch_cfgs': [[i, end_mapping_ids[i] + 1]],
+                    'stitch_layer_ids': [i // (depths[1] // depths[0])]
+                })
+
+
+    # large --> small --> large
+    for ls_cfg in ls_configs:
+        for sl_cfg in sl_configs:
+            if sl_cfg['stitch_layer_ids'][0] == depths[0] - 1:
+                continue
+            if sl_cfg['stitch_cfgs'][0][0] >= ls_cfg['stitch_cfgs'][0][1]:
+                lsl_confgs.append({
+                    'comb_id': [1, 0, 1],
+                    'stitch_cfgs': [ls_cfg['stitch_cfgs'][0], sl_cfg['stitch_cfgs'][0]],
+                    'stitch_layer_ids': ls_cfg['stitch_layer_ids'] + sl_cfg['stitch_layer_ids']
+                })
+
+    # small --> large --> small
+    sls_configs = []
+    for sl_cfg in sl_configs:
+        for ls_cfg in ls_configs:
+            if ls_cfg['stitch_cfgs'][0][0] >= sl_cfg['stitch_cfgs'][0][1]:
+                sls_configs.append({
+                    'comb_id': [0, 1, 0],
+                    'stitch_cfgs': [sl_cfg['stitch_cfgs'][0], ls_cfg['stitch_cfgs'][0]],
+                    'stitch_layer_ids': sl_cfg['stitch_layer_ids'] + ls_cfg['stitch_layer_ids']
+                })
+
+    total_configs += sl_configs + ls_configs + lsl_confgs + sls_configs
+
+    anchor_ids = []
+    sl_ids = []
+    ls_ids = []
+    lsl_ids = []
+    sls_ids = []
+
+    for i, cfg in enumerate(total_configs):
+        comb_id = cfg['comb_id']
+
+        if len(comb_id) == 1:
+            anchor_ids.append(i)
+            continue
+
+        if len(comb_id) == 2:
+            route = []
+            front, end = cfg['stitch_cfgs'][0]
+            route.append([0, front])
+            route.append([end, depths[comb_id[-1]]])
+            cfg['route'] = route
+            
+            if comb_id == [0, 1] and front != 11:
+                sl_ids.append(i)
+            elif comb_id == [1, 0]:
+                ls_ids.append(i)
+
+        if len(comb_id) == 3:
+            route = []
+            front_1, end_1 = cfg['stitch_cfgs'][0]
+            front_2, end_2 = cfg['stitch_cfgs'][1]
+            route.append([0, front_1])
+            route.append([end_1, front_2])
+            route.append([end_2, depths[comb_id[-1]]])
+            cfg['route'] = route
+
+            if comb_id == [1, 0, 1]:
+                lsl_ids.append(i)
+            elif comb_id == [0, 1, 0]:
+                sls_ids.append(i)
+
+        cfg['stitch_layer_ids'].append(-1)
+
+    model_combos = [(0, 1), (1, 0)]
+    return total_configs, model_combos, [len(sl_configs), len(ls_configs)], anchor_ids, sl_ids, ls_ids, lsl_ids, sls_ids
+
+
+def format_out_features(outs, with_cls_token, hw_shape):
+    B, _, C = outs[0].shape
+    for i in range(len(outs)):
+        if with_cls_token:
+            # Remove class token and reshape token for decoder head
+            outs[i] = outs[i][:, 1:].reshape(B, hw_shape[0], hw_shape[1],
+                                             C).permute(0, 3, 1, 2).contiguous()
+        else:
+            outs[i] = outs[i].reshape(B, hw_shape[0], hw_shape[1],
+                                      C).permute(0, 3, 1, 2).contiguous()
+    return outs
+
+
+class LoRALayer():
+    def __init__(
+        self,
+        r: int,
+        lora_alpha: int,
+        lora_dropout: float,
+        merge_weights: bool,
+    ):
+        self.r = r
+        self.lora_alpha = lora_alpha
+        # Optional dropout
+        if lora_dropout > 0.:
+            self.lora_dropout = nn.Dropout(p=lora_dropout)
+        else:
+            self.lora_dropout = lambda x: x
+        # Mark the weight as unmerged
+        self.merged = False
+        self.merge_weights = merge_weights
+
+class Linear(nn.Linear, LoRALayer):
+    # LoRA implemented in a dense layer
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        r: int = 0,
+        lora_alpha: int = 1,
+        lora_dropout: float = 0.,
+        fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
+        merge_weights: bool = True,
+        **kwargs
+    ):
+        nn.Linear.__init__(self, in_features, out_features, **kwargs)
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
+                           merge_weights=merge_weights)
+
+        self.fan_in_fan_out = fan_in_fan_out
+        # Actual trainable parameters
+        if r > 0:
+            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
+            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))
+            self.scaling = self.lora_alpha / self.r
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+        self.reset_parameters()
+        if fan_in_fan_out:
+            self.weight.data = self.weight.data.transpose(0, 1)
+
+    def reset_parameters(self):
+        nn.Linear.reset_parameters(self)
+        if hasattr(self, 'lora_A'):
+            # initialize A the same way as the default for nn.Linear and B to zero
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            nn.init.zeros_(self.lora_B)
+
+    def train(self, mode: bool = True):
+        def T(w):
+            return w.transpose(0, 1) if self.fan_in_fan_out else w
+        nn.Linear.train(self, mode)
+        if mode:
+            if self.merge_weights and self.merged:
+                # Make sure that the weights are not merged
+                if self.r > 0:
+                    self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
+                self.merged = False
+        else:
+            if self.merge_weights and not self.merged:
+                # Merge the weights and mark it
+                if self.r > 0:
+                    self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
+                self.merged = True
+
+    def forward(self, x: torch.Tensor):
+        def T(w):
+            return w.transpose(0, 1) if self.fan_in_fan_out else w
+        if self.r > 0 and not self.merged:
+            result = F.linear(x, T(self.weight), bias=self.bias)
+            if self.r > 0:
+                result += (self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) * self.scaling
+            return result
+        else:
+            return F.linear(x, T(self.weight), bias=self.bias)
+
+
+class StitchingLayer(nn.Module):
+    def __init__(self, in_features=None, out_features=None, r=0):
+        super().__init__()
+        self.transform = Linear(in_features, out_features,  r=r)
+
+    def init_stitch_weights_bias(self, weight, bias):
+        self.transform.weight.data.copy_(weight)
+        self.transform.bias.data.copy_(bias)
+
+    def forward(self, x):
+        out = self.transform(x)
+        return out
+
+
+class SNNet(nn.Module):
+
+    def __init__(self, anchors=None):
+        super(SNNet, self).__init__()
+        self.anchors = nn.ModuleList(anchors)
+
+        self.depths = [len(anc.blocks) for anc in self.anchors]
+
+        total_configs, num_stitches = get_stitch_configs_general_unequal(self.depths)
+        self.stitch_layers = nn.ModuleList(
+            [StitchingLayer(self.anchors[0].embed_dim, self.anchors[1].embed_dim) for _ in range(num_stitches)])
+
+        self.stitch_configs = {i: cfg for i, cfg in enumerate(total_configs)}
+        self.all_cfgs = list(self.stitch_configs.keys())
+        self.num_configs = len(self.all_cfgs)
+        self.stitch_config_id = 0
+        self.is_ranking = False
+
+    def reset_stitch_id(self, stitch_config_id):
+        self.stitch_config_id = stitch_config_id
+
+    def initialize_stitching_weights(self, x):
+        logger = get_root_logger()
+        front, end = 0, 1
+        with torch.no_grad():
+            front_features = self.anchors[front].extract_block_features(x)
+            end_features = self.anchors[end].extract_block_features(x)
+
+        for i, blk_id in enumerate(range(self.depths[0])):
+            front_id, end_id = i, (i + 1) * (self.depths[1] // self.depths[0])
+            front_blk_feat = front_features[front_id]
+            end_blk_feat = end_features[end_id - 1]
+            w, b = ps_inv(front_blk_feat, end_blk_feat)
+            self.stitch_layers[i].init_stitch_weights_bias(w, b)
+            logger.info(f'Initialized Stitching Model {front} to Model {end}, Layer {i}')
+
+    def init_weights(self):
+        for anc in self.anchors:
+            anc.init_weights()
+
+    def sampling_stitch_config(self):
+        self.stitch_config_id = np.random.choice(self.all_cfgs)
+
+    def forward(self, x):
+
+        stitch_cfg_id = self.stitch_config_id
+        comb_id = self.stitch_configs[stitch_cfg_id]['comb_id']
+
+        if len(comb_id) == 1:
+            return self.anchors[comb_id[0]](x)
+
+        cfg = self.stitch_configs[stitch_cfg_id]['stitch_cfgs']
+
+        x = self.anchors[comb_id[0]].forward_until(x, blk_id=cfg[0])
+        x = self.stitch_layers[cfg[0]](x)
+        x = self.anchors[comb_id[1]].forward_from(x, blk_id=cfg[1])
+
+        return x
+
+
+class SNNetv2(nn.Module):
+
+    def __init__(self, anchors=None, include_sl=True, include_ls=True, include_lsl=True, include_sls=True, lora_r=0):
+        super(SNNetv2, self).__init__()
+        self.anchors = nn.ModuleList(anchors)
+
+        self.lora_r = lora_r
+
+        self.depths = [len(anc.blocks) for anc in self.anchors]
+
+        total_configs, model_combos, num_stitches, anchor_ids, sl_ids, ls_ids, lsl_ids, sls_ids = get_stitch_configs_bidirection(self.depths)
+
+        self.stitch_layers = nn.ModuleList()
+        self.stitching_map_id = {}
+
+        for i, (comb, num_sth) in enumerate(zip(model_combos, num_stitches)):
+            front, end = comb
+            temp = nn.ModuleList(
+                [StitchingLayer(self.anchors[front].embed_dim, self.anchors[end].embed_dim, r=lora_r) for _ in range(num_sth)])
+            temp.append(nn.Identity())
+            self.stitch_layers.append(temp)
+
+        self.stitch_configs = {i: cfg for i, cfg in enumerate(total_configs)}
+        self.stitch_init_configs = {i: cfg for i, cfg in enumerate(total_configs) if len(cfg['comb_id']) == 2}
+
+
+        self.all_cfgs = list(self.stitch_configs.keys())
+        logger = get_root_logger()
+        logger.info(str(self.all_cfgs))
+
+
+        self.all_cfgs = anchor_ids
+
+        if include_sl:
+            self.all_cfgs += sl_ids
+
+        if include_ls:
+            self.all_cfgs += ls_ids
+
+        if include_lsl:
+            self.all_cfgs += lsl_ids
+
+        if include_sls:
+            self.all_cfgs += sls_ids
+
+        self.num_configs = len(self.stitch_configs)
+        self.stitch_config_id = 0
+
+    def reset_stitch_id(self, stitch_config_id):
+        self.stitch_config_id = stitch_config_id
+
+    def set_ranking_mode(self, ranking_mode):
+        self.is_ranking = ranking_mode
+
+    def initialize_stitching_weights(self, x):
+        logger = get_root_logger()
+        anchor_features = []
+        for anchor in self.anchors:
+            with torch.no_grad():
+                temp = anchor.extract_block_features(x)
+                anchor_features.append(temp)
+
+        for idx, cfg in self.stitch_init_configs.items():
+            comb_id = cfg['comb_id']
+            if len(comb_id) == 2:
+                front_id, end_id = cfg['stitch_cfgs'][0]
+                stitch_layer_id = cfg['stitch_layer_ids'][0]
+                front_blk_feat = anchor_features[comb_id[0]][front_id]
+                end_blk_feat = anchor_features[comb_id[1]][end_id - 1]
+                w, b = ps_inv(front_blk_feat, end_blk_feat)
+                self.stitch_layers[comb_id[0]][stitch_layer_id].init_stitch_weights_bias(w, b)
+                logger.info(f'Initialized Stitching Layer {cfg}')
+
+    def init_weights(self):
+        for anc in self.anchors:
+            anc.init_weights()
+
+
+    def sampling_stitch_config(self):
+        flops_id = np.random.choice(len(self.flops_grouped_cfgs), p=self.flops_sampling_probs)
+        stitch_config_id = np.random.choice(self.flops_grouped_cfgs[flops_id])
+        return stitch_config_id
+
+    def forward(self, x):
+
+        if self.training:
+            stitch_cfg_id = self.sampling_stitch_config()
+        else:
+            stitch_cfg_id = self.stitch_config_id
+
+        comb_id = self.stitch_configs[stitch_cfg_id]['comb_id']
+
+        # forward by a single anchor
+        if len(comb_id) == 1:
+            return self.anchors[comb_id[0]](x)
+
+        # forward among anchors
+        route = self.stitch_configs[stitch_cfg_id]['route']
+        stitch_layer_ids = self.stitch_configs[stitch_cfg_id]['stitch_layer_ids']
+
+        # patch embeding
+        x = self.anchors[comb_id[0]].forward_patch_embed(x)
+
+        for i, (model_id, cfg) in enumerate(zip(comb_id, route)):
+
+            x = self.anchors[model_id].selective_forward(x, cfg[0], cfg[1])
+            x = self.stitch_layers[model_id][stitch_layer_ids[i]](x)
+
+        x = self.anchors[comb_id[-1]].forward_norm_head(x)
+        return x
+

snnetv2_deit3_s_l.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d455f17d73f4ed74702076d4cea516194d8c4aa8fbbc63192f85795f79c76b4
+size 1350494458

stitches_res_s_l.txt

@@ -0,0 +1,134 @@
+{"loss": 0.7156664722345092, "acc1": 82.9060024609375, "acc5": 96.73400244140625, "cfg_id": 0, "flops": 4608338304}
+{"loss": 0.5377805712209507, "acc1": 86.97800256835937, "acc5": 98.2540023046875, "cfg_id": 1, "flops": 61604135936}
+{"loss": 0.5598483879796483, "acc1": 86.57800241210937, "acc5": 98.08200240234375, "cfg_id": 2, "flops": 56843745792}
+{"loss": 0.5534007405354218, "acc1": 86.6480025390625, "acc5": 98.1760025390625, "cfg_id": 3, "flops": 52102230016}
+{"loss": 0.5610568577028585, "acc1": 86.49800245117187, "acc5": 98.06600229492187, "cfg_id": 4, "flops": 47360714240}
+{"loss": 0.5747850706067049, "acc1": 86.26000259765625, "acc5": 97.93800240234376, "cfg_id": 5, "flops": 42619198464}
+{"loss": 0.5890085864812136, "acc1": 85.79200244140625, "acc5": 97.80000272460937, "cfg_id": 6, "flops": 37877682688}
+{"loss": 0.6165087098876635, "acc1": 85.08200264648437, "acc5": 97.55600231445312, "cfg_id": 7, "flops": 33136166912}
+{"loss": 0.6652509210574807, "acc1": 83.69200263671875, "acc5": 97.23600259765625, "cfg_id": 8, "flops": 28394651136}
+{"loss": 0.7374675334290122, "acc1": 81.7120026171875, "acc5": 96.53200251953125, "cfg_id": 9, "flops": 23653135360}
+{"loss": 0.7991558508665273, "acc1": 79.50600241210938, "acc5": 95.90200240234375, "cfg_id": 10, "flops": 18911619584}
+{"loss": 0.7554851990531791, "acc1": 80.63600265625, "acc5": 96.09000245117187, "cfg_id": 11, "flops": 14170103808}
+{"loss": 0.7068120487824534, "acc1": 82.25000237304687, "acc5": 96.35600284179688, "cfg_id": 12, "flops": 9428588032}
+{"loss": 0.7329587066038088, "acc1": 82.6600027734375, "acc5": 96.58200264648437, "cfg_id": 14, "flops": 9523655552}
+{"loss": 0.7238117807516546, "acc1": 82.94800252929687, "acc5": 96.68600232421875, "cfg_id": 15, "flops": 14265171328}
+{"loss": 0.7139950410434694, "acc1": 83.0860026953125, "acc5": 96.75800252929687, "cfg_id": 16, "flops": 19006687104}
+{"loss": 0.7004092067028537, "acc1": 83.25400249023437, "acc5": 96.8740026171875, "cfg_id": 17, "flops": 23748202880}
+{"loss": 0.6828147762201049, "acc1": 83.45000244140626, "acc5": 96.9520026171875, "cfg_id": 18, "flops": 28489718656}
+{"loss": 0.6787144099221085, "acc1": 83.56400258789063, "acc5": 97.0600024609375, "cfg_id": 19, "flops": 33231234432}
+{"loss": 0.6765228407175252, "acc1": 83.43400251953125, "acc5": 97.19200266601563, "cfg_id": 20, "flops": 37972750208}
+{"loss": 0.6841061733888857, "acc1": 83.5900022265625, "acc5": 97.20800275390626, "cfg_id": 21, "flops": 42714265984}
+{"loss": 0.6446758140104286, "acc1": 84.8660023828125, "acc5": 97.44400258789062, "cfg_id": 22, "flops": 47455781760}
+{"loss": 0.5939652780917558, "acc1": 86.23000265625, "acc5": 97.69200270507812, "cfg_id": 23, "flops": 52197297536}
+{"loss": 0.5654762382760192, "acc1": 86.43400250976562, "acc5": 97.632002578125, "cfg_id": 24, "flops": 56938813312}
+{"loss": 0.5636055112788172, "acc1": 86.39000270507813, "acc5": 98.04800252929688, "cfg_id": 25, "flops": 57017547264}
+{"loss": 0.5706944450397383, "acc1": 86.234002578125, "acc5": 98.00000237304687, "cfg_id": 26, "flops": 52276031488}
+{"loss": 0.5833309799658529, "acc1": 85.9240025390625, "acc5": 97.9260024609375, "cfg_id": 27, "flops": 47534515712}
+{"loss": 0.5972222860225223, "acc1": 85.57400262695313, "acc5": 97.70800255859375, "cfg_id": 28, "flops": 42792999936}
+{"loss": 0.6253456006560362, "acc1": 84.89800259765624, "acc5": 97.47400255859375, "cfg_id": 29, "flops": 38051484160}
+{"loss": 0.6745385262889393, "acc1": 83.5380026171875, "acc5": 97.07600244140625, "cfg_id": 30, "flops": 33309968384}
+{"loss": 0.7486309014034994, "acc1": 81.42600245117187, "acc5": 96.33600258789062, "cfg_id": 31, "flops": 28568452608}
+{"loss": 0.8134756960877867, "acc1": 79.16000235351562, "acc5": 95.72400271484375, "cfg_id": 32, "flops": 23826936832}
+{"loss": 0.7671100513050051, "acc1": 80.37200240234375, "acc5": 95.98400258789063, "cfg_id": 33, "flops": 19085421056}
+{"loss": 0.7206548866674756, "acc1": 81.91000239257812, "acc5": 96.23800239257812, "cfg_id": 34, "flops": 14343905280}
+{"loss": 0.5626872230998494, "acc1": 86.44600235351562, "acc5": 98.062002421875, "cfg_id": 35, "flops": 57017547264}
+{"loss": 0.5785287711769342, "acc1": 86.06400251953124, "acc5": 97.9420023046875, "cfg_id": 36, "flops": 52276031488}
+{"loss": 0.5930487287202568, "acc1": 85.78400234375, "acc5": 97.79000255859376, "cfg_id": 37, "flops": 47534515712}
+{"loss": 0.6189901619923838, "acc1": 85.10800268554688, "acc5": 97.50400228515625, "cfg_id": 38, "flops": 42792999936}
+{"loss": 0.6674688318462083, "acc1": 83.76600272460938, "acc5": 97.09400264648437, "cfg_id": 39, "flops": 38051484160}
+{"loss": 0.7388352820593299, "acc1": 81.70200266601563, "acc5": 96.47000241210938, "cfg_id": 40, "flops": 33309968384}
+{"loss": 0.803126322613521, "acc1": 79.4560025390625, "acc5": 95.81400245117187, "cfg_id": 41, "flops": 28568452608}
+{"loss": 0.7581946616145697, "acc1": 80.70800243164062, "acc5": 96.08600255859375, "cfg_id": 42, "flops": 23826936832}
+{"loss": 0.7118472667467414, "acc1": 82.22600248046875, "acc5": 96.31000268554688, "cfg_id": 43, "flops": 19085421056}
+{"loss": 0.5727639499713074, "acc1": 86.2180025, "acc5": 97.98200247070312, "cfg_id": 44, "flops": 57017547264}
+{"loss": 0.5866389607615543, "acc1": 85.84400263671876, "acc5": 97.8600024609375, "cfg_id": 45, "flops": 52276031488}
+{"loss": 0.6107792718279542, "acc1": 85.19800255859376, "acc5": 97.61800235351562, "cfg_id": 46, "flops": 47534515712}
+{"loss": 0.6602028349809574, "acc1": 83.92600266601562, "acc5": 97.23800282226563, "cfg_id": 47, "flops": 42792999936}
+{"loss": 0.7285334389431007, "acc1": 82.0040028125, "acc5": 96.52400247070312, "cfg_id": 48, "flops": 38051484160}
+{"loss": 0.7910783413910505, "acc1": 79.69600262695313, "acc5": 95.95800241210938, "cfg_id": 49, "flops": 33309968384}
+{"loss": 0.7478298004152197, "acc1": 80.89400243164063, "acc5": 96.172002421875, "cfg_id": 50, "flops": 28568452608}
+{"loss": 0.7014034044449077, "acc1": 82.45600264648438, "acc5": 96.438002421875, "cfg_id": 51, "flops": 23826936832}
+{"loss": 0.5799332931637764, "acc1": 85.92400239257813, "acc5": 97.94000249023438, "cfg_id": 52, "flops": 57017547264}
+{"loss": 0.6004864230300441, "acc1": 85.43800255859375, "acc5": 97.70800227539063, "cfg_id": 53, "flops": 52276031488}
+{"loss": 0.647012604287628, "acc1": 84.20200264648437, "acc5": 97.30600264648437, "cfg_id": 54, "flops": 47534515712}
+{"loss": 0.7162722434961435, "acc1": 82.29000248046874, "acc5": 96.6640023046875, "cfg_id": 55, "flops": 42792999936}
+{"loss": 0.7757266998065241, "acc1": 79.9760025, "acc5": 96.050002421875, "cfg_id": 56, "flops": 38051484160}
+{"loss": 0.7351311787285588, "acc1": 81.04400232421875, "acc5": 96.2400026953125, "cfg_id": 57, "flops": 33309968384}
+{"loss": 0.6896895027408997, "acc1": 82.6220026171875, "acc5": 96.55400252929688, "cfg_id": 58, "flops": 28568452608}
+{"loss": 0.5911911701727094, "acc1": 85.53000266601562, "acc5": 97.76600241210937, "cfg_id": 59, "flops": 57017547264}
+{"loss": 0.6371258264125297, "acc1": 84.41200249023437, "acc5": 97.44200245117187, "cfg_id": 60, "flops": 52276031488}
+{"loss": 0.7022040815403064, "acc1": 82.49400240234375, "acc5": 96.74400272460937, "cfg_id": 61, "flops": 47534515712}
+{"loss": 0.7612808859257987, "acc1": 80.29200265625, "acc5": 96.15800239257813, "cfg_id": 62, "flops": 42792999936}
+{"loss": 0.7246641330420971, "acc1": 81.20400250976563, "acc5": 96.42400264648437, "cfg_id": 63, "flops": 38051484160}
+{"loss": 0.6782861414619468, "acc1": 82.8040024609375, "acc5": 96.60800270507812, "cfg_id": 64, "flops": 33309968384}
+{"loss": 0.629801401755575, "acc1": 84.65200262695312, "acc5": 97.54200265625, "cfg_id": 65, "flops": 57017547264}
+{"loss": 0.6992729283643492, "acc1": 82.58200259765626, "acc5": 96.85600262695313, "cfg_id": 66, "flops": 52276031488}
+{"loss": 0.7595290538262237, "acc1": 80.35600247070313, "acc5": 96.27000247070312, "cfg_id": 67, "flops": 47534515712}
+{"loss": 0.7238247728709019, "acc1": 81.37600248046876, "acc5": 96.46200267578125, "cfg_id": 68, "flops": 42792999936}
+{"loss": 0.6760879844765771, "acc1": 82.96800264648438, "acc5": 96.69400274414062, "cfg_id": 69, "flops": 38051484160}
+{"loss": 0.68392569430624, "acc1": 83.16200254882813, "acc5": 97.09200258789062, "cfg_id": 70, "flops": 57017547264}
+{"loss": 0.7509645553249301, "acc1": 80.68000260742187, "acc5": 96.3900025, "cfg_id": 71, "flops": 52276031488}
+{"loss": 0.7208586267449639, "acc1": 81.55200274414062, "acc5": 96.62200272460937, "cfg_id": 72, "flops": 47534515712}
+{"loss": 0.6785354860352747, "acc1": 82.86000262695312, "acc5": 96.80000244140625, "cfg_id": 73, "flops": 42792999936}
+{"loss": 0.7184764705598354, "acc1": 81.61200241210938, "acc5": 96.63200258789063, "cfg_id": 74, "flops": 57017547264}
+{"loss": 0.7229886900520686, "acc1": 81.45000249023437, "acc5": 96.62200250976562, "cfg_id": 75, "flops": 52276031488}
+{"loss": 0.6883746685855316, "acc1": 83.01600272460938, "acc5": 96.83200240234375, "cfg_id": 76, "flops": 47534515712}
+{"loss": 0.6293963799535325, "acc1": 83.90800231445313, "acc5": 97.27400245117188, "cfg_id": 77, "flops": 57017547264}
+{"loss": 0.642419446824175, "acc1": 84.31400258789063, "acc5": 97.19200287109375, "cfg_id": 78, "flops": 52276031488}
+{"loss": 0.5880116202275861, "acc1": 85.90600231445312, "acc5": 97.58400263671875, "cfg_id": 79, "flops": 57017547264}
+{"loss": 0.750676692096573, "acc1": 82.14200271484376, "acc5": 96.36600260742188, "cfg_id": 80, "flops": 9504781184}
+{"loss": 0.7431871895537232, "acc1": 82.234002578125, "acc5": 96.39200241210938, "cfg_id": 81, "flops": 14246296960}
+{"loss": 0.7236298957105839, "acc1": 82.62600249023437, "acc5": 96.57600243164063, "cfg_id": 82, "flops": 18987812736}
+{"loss": 0.7074674766397837, "acc1": 82.84600237304687, "acc5": 96.68800247070313, "cfg_id": 83, "flops": 23729328512}
+{"loss": 0.7014015182062532, "acc1": 82.99000265625, "acc5": 96.7740025, "cfg_id": 84, "flops": 28470844288}
+{"loss": 0.6996880258348855, "acc1": 82.98000252929687, "acc5": 96.90200263671875, "cfg_id": 85, "flops": 33212360064}
+{"loss": 0.7077699161953095, "acc1": 82.96200270507812, "acc5": 96.98600258789062, "cfg_id": 86, "flops": 37953875840}
+{"loss": 0.6674120087515224, "acc1": 84.33000274414063, "acc5": 97.2640025, "cfg_id": 87, "flops": 42695391616}
+{"loss": 0.6169534720141779, "acc1": 85.86200280273438, "acc5": 97.49200272460938, "cfg_id": 88, "flops": 47436907392}
+{"loss": 0.5848360503600403, "acc1": 86.02600271484376, "acc5": 97.4480026171875, "cfg_id": 89, "flops": 52178423168}
+{"loss": 0.7346750153510859, "acc1": 82.5540027734375, "acc5": 96.52400241210937, "cfg_id": 90, "flops": 9504781184}
+{"loss": 0.7158081559182117, "acc1": 82.82200255859375, "acc5": 96.65600255859376, "cfg_id": 91, "flops": 14246296960}
+{"loss": 0.6994372372600165, "acc1": 83.03600239257813, "acc5": 96.74600252929687, "cfg_id": 92, "flops": 18987812736}
+{"loss": 0.6947964186582601, "acc1": 83.07400255859375, "acc5": 96.88200241210937, "cfg_id": 93, "flops": 23729328512}
+{"loss": 0.6946824553112189, "acc1": 83.0200026171875, "acc5": 96.99400251953125, "cfg_id": 94, "flops": 28470844288}
+{"loss": 0.7041463901599249, "acc1": 83.21600236328125, "acc5": 97.00400250976563, "cfg_id": 95, "flops": 33212360064}
+{"loss": 0.6699620116163384, "acc1": 84.54600258789063, "acc5": 97.31600244140625, "cfg_id": 96, "flops": 37953875840}
+{"loss": 0.6176637105192199, "acc1": 85.91800228515625, "acc5": 97.57000255859376, "cfg_id": 97, "flops": 42695391616}
+{"loss": 0.5765587539045196, "acc1": 86.1880023046875, "acc5": 97.54000249023437, "cfg_id": 98, "flops": 47436907392}
+{"loss": 0.7319535712401072, "acc1": 82.49200258789062, "acc5": 96.46600271484375, "cfg_id": 99, "flops": 9504781184}
+{"loss": 0.710809505516381, "acc1": 82.7860023046875, "acc5": 96.6260026171875, "cfg_id": 100, "flops": 14246296960}
+{"loss": 0.7044268037107858, "acc1": 82.93000239257813, "acc5": 96.79800258789062, "cfg_id": 101, "flops": 18987812736}
+{"loss": 0.7076575808001287, "acc1": 82.83200243164063, "acc5": 96.8820025390625, "cfg_id": 102, "flops": 23729328512}
+{"loss": 0.7188302328189214, "acc1": 82.88200259765625, "acc5": 96.93600249023437, "cfg_id": 103, "flops": 28470844288}
+{"loss": 0.6856357377361167, "acc1": 84.2520023046875, "acc5": 97.2200026171875, "cfg_id": 104, "flops": 33212360064}
+{"loss": 0.6273381847210906, "acc1": 85.758002734375, "acc5": 97.44800275390625, "cfg_id": 105, "flops": 37953875840}
+{"loss": 0.5830204013847944, "acc1": 86.01000260742188, "acc5": 97.4780026171875, "cfg_id": 106, "flops": 42695391616}
+{"loss": 0.7305513945492831, "acc1": 82.3180023828125, "acc5": 96.47200256835937, "cfg_id": 107, "flops": 9504781184}
+{"loss": 0.7206297208639708, "acc1": 82.37200228515626, "acc5": 96.61600234375, "cfg_id": 108, "flops": 14246296960}
+{"loss": 0.7241401795975186, "acc1": 82.33800244140625, "acc5": 96.74400267578125, "cfg_id": 109, "flops": 18987812736}
+{"loss": 0.7350799917723193, "acc1": 82.57800231445313, "acc5": 96.78000252929688, "cfg_id": 110, "flops": 23729328512}
+{"loss": 0.7013292148935072, "acc1": 83.95000255859375, "acc5": 97.17600254882812, "cfg_id": 111, "flops": 28470844288}
+{"loss": 0.64130035031474, "acc1": 85.54600244140624, "acc5": 97.36800262695313, "cfg_id": 112, "flops": 33212360064}
+{"loss": 0.5961506485826138, "acc1": 85.76800252929688, "acc5": 97.33400268554688, "cfg_id": 113, "flops": 37953875840}
+{"loss": 0.7443677056580782, "acc1": 81.86200244140625, "acc5": 96.29200241210937, "cfg_id": 114, "flops": 9504781184}
+{"loss": 0.7442678388659701, "acc1": 81.82000262695313, "acc5": 96.3900026953125, "cfg_id": 115, "flops": 14246296960}
+{"loss": 0.749958168037913, "acc1": 81.99400229492187, "acc5": 96.49600264648437, "cfg_id": 116, "flops": 18987812736}
+{"loss": 0.7073916116672935, "acc1": 83.43200267578125, "acc5": 96.99200250976563, "cfg_id": 117, "flops": 23729328512}
+{"loss": 0.6501240834706661, "acc1": 85.1480025, "acc5": 97.212002734375, "cfg_id": 118, "flops": 28470844288}
+{"loss": 0.6135486943477934, "acc1": 85.35400266601563, "acc5": 97.16000249023438, "cfg_id": 119, "flops": 33212360064}
+{"loss": 0.7824224890633062, "acc1": 81.16600250976562, "acc5": 96.11200255859374, "cfg_id": 120, "flops": 9504781184}
+{"loss": 0.7932735298844901, "acc1": 80.93600233398438, "acc5": 96.12600254882813, "cfg_id": 121, "flops": 14246296960}
+{"loss": 0.7476008046757091, "acc1": 82.49800259765625, "acc5": 96.65000265625, "cfg_id": 122, "flops": 18987812736}
+{"loss": 0.6842290776019747, "acc1": 84.35800275390625, "acc5": 96.94000270507813, "cfg_id": 123, "flops": 23729328512}
+{"loss": 0.6409159135073423, "acc1": 84.520002578125, "acc5": 96.91800252929687, "cfg_id": 124, "flops": 28470844288}
+{"loss": 0.8394820786109476, "acc1": 79.96600270507813, "acc5": 95.5980024609375, "cfg_id": 125, "flops": 9504781184}
+{"loss": 0.7924092794683847, "acc1": 81.1100024609375, "acc5": 96.12600264648438, "cfg_id": 126, "flops": 14246296960}
+{"loss": 0.724013783997207, "acc1": 83.04000248046874, "acc5": 96.46600255859374, "cfg_id": 127, "flops": 18987812736}
+{"loss": 0.6948224610338608, "acc1": 83.016002421875, "acc5": 96.44000282226563, "cfg_id": 128, "flops": 23729328512}
+{"loss": 0.8688964597655066, "acc1": 79.15800264160156, "acc5": 95.2460025, "cfg_id": 129, "flops": 9504781184}
+{"loss": 0.8095247168658357, "acc1": 80.57800265625, "acc5": 95.63000274414063, "cfg_id": 130, "flops": 14246296960}
+{"loss": 0.7750140779059042, "acc1": 80.9680026171875, "acc5": 95.68200263671875, "cfg_id": 131, "flops": 18987812736}
+{"loss": 0.9017180648039688, "acc1": 77.88400251953125, "acc5": 94.78400234375, "cfg_id": 132, "flops": 9504781184}
+{"loss": 0.8799216277671583, "acc1": 77.84800252929688, "acc5": 94.79800247070312, "cfg_id": 133, "flops": 14246296960}
+{"loss": 0.8987371790589709, "acc1": 78.12000266601562, "acc5": 94.92200256835937, "cfg_id": 134, "flops": 9504781184}

utils.py

@@ -0,0 +1,408 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+"""
+Misc functions, including distributed helpers.
+
+Mostly copy-paste from torchvision references.
+"""
+import io
+import os
+import time
+from collections import defaultdict, deque
+import datetime
+
+import torch
+import torch.distributed as dist
+import logging
+
+logger_initialized = {}
+
+def group_subnets_by_flops(data, flops_gap=1.0):
+    sorted_data = {k: v for k, v in sorted(data.items(), key=lambda item: item[1])}
+    candidate_idx = []
+    grouped_cands = []
+    last_flops = 0
+    for cfg_id, flops in sorted_data.items():
+        flops = flops / 1e9
+        if abs(last_flops - flops) > flops_gap:
+            if len(candidate_idx) > 0:
+                grouped_cands.append(sorted(candidate_idx))
+            candidate_idx = [int(cfg_id)]
+            last_flops = flops
+        else:
+            candidate_idx.append(int(cfg_id))
+
+    if len(candidate_idx) > 0:
+        grouped_cands.append(sorted(candidate_idx))
+
+    return grouped_cands
+
+def find_best_candidates(data):
+    sorted_data = {k: v for k, v in sorted(data.items(), key=lambda item: item[1])}
+    candidate_idx = []
+    last_flops = 0
+    for cfg_id, values in sorted_data.items():
+        flops, score = values
+        if abs(last_flops - flops) > 1:
+            candidate_idx.append(cfg_id)
+            last_flops = flops
+        else:
+            if score > data[candidate_idx[-1]][1]:
+                candidate_idx[-1] = cfg_id
+
+    return candidate_idx
+
+
+
+def find_top_candidates(data, ratio=0.9):
+    sorted_data = {k: v for k, v in sorted(data.items(), key=lambda item: item[1])}
+    candidate_idx = []
+    grouped_cands = []
+    last_flops = 0
+    for cfg_id, values in sorted_data.items():
+        flops, score = values
+        if abs(last_flops - flops) > 3:
+            if len(candidate_idx) > 0:
+                grouped_cands.append(candidate_idx)
+            candidate_idx = [cfg_id]
+            last_flops = flops
+        else:
+            candidate_idx.append(cfg_id)
+
+    if len(candidate_idx) > 0:
+        grouped_cands.append(candidate_idx)
+
+    final_list = []
+    for group in grouped_cands:
+        if len(group) == 1:
+            final_list += list(map(int, group))
+            continue
+        scores = torch.tensor([sorted_data[cfg_id][-1] for cfg_id in group])
+
+        indices = torch.argsort(scores, descending=True)
+        num_selected = int(ratio*len(group)) if int(ratio*len(group)) > 0 else 1
+
+        top_ids = indices[:num_selected].tolist()
+        selected = [group[idx] for idx in top_ids]
+        final_list += list(map(int, selected))
+
+    return final_list
+
+
+
+def get_logger(name, log_file=None, log_level=logging.INFO, file_mode='w'):
+    """Initialize and get a logger by name.
+
+    If the logger has not been initialized, this method will initialize the
+    logger by adding one or two handlers, otherwise the initialized logger will
+    be directly returned. During initialization, a StreamHandler will always be
+    added. If `log_file` is specified and the process rank is 0, a FileHandler
+    will also be added.
+
+    Args:
+        name (str): Logger name.
+        log_file (str | None): The log filename. If specified, a FileHandler
+            will be added to the logger.
+        log_level (int): The logger level. Note that only the process of
+            rank 0 is affected, and other processes will set the level to
+            "Error" thus be silent most of the time.
+        file_mode (str): The file mode used in opening log file.
+            Defaults to 'w'.
+
+    Returns:
+        logging.Logger: The expected logger.
+    """
+    logger = logging.getLogger(name)
+    if name in logger_initialized:
+        return logger
+    # handle hierarchical names
+    # e.g., logger "a" is initialized, then logger "a.b" will skip the
+    # initialization since it is a child of "a".
+    for logger_name in logger_initialized:
+        if name.startswith(logger_name):
+            return logger
+
+    stream_handler = logging.StreamHandler()
+    handlers = [stream_handler]
+
+    if dist.is_available() and dist.is_initialized():
+        rank = dist.get_rank()
+    else:
+        rank = 0
+
+    # only rank 0 will add a FileHandler
+    if rank == 0 and log_file is not None:
+        # Here, the default behaviour of the official logger is 'a'. Thus, we
+        # provide an interface to change the file mode to the default
+        # behaviour.
+        file_handler = logging.FileHandler(log_file, file_mode)
+        handlers.append(file_handler)
+
+    formatter = logging.Formatter(
+        '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    for handler in handlers:
+        handler.setFormatter(formatter)
+        handler.setLevel(log_level)
+        logger.addHandler(handler)
+
+    if rank == 0:
+        logger.setLevel(log_level)
+    else:
+        logger.setLevel(logging.ERROR)
+
+    logger_initialized[name] = True
+
+    return logger
+
+
+def get_root_logger(log_file=None, log_level=logging.INFO):
+    """Get the root logger.
+
+    The logger will be initialized if it has not been initialized. By default a
+    StreamHandler will be added. If `log_file` is specified, a FileHandler will
+    also be added. The name of the root logger is the top-level package name,
+    e.g., "mmseg".
+
+    Args:
+        log_file (str | None): The log filename. If specified, a FileHandler
+            will be added to the root logger.
+        log_level (int): The root logger level. Note that only the process of
+            rank 0 is affected, while other processes will set the level to
+            "Error" and be silent most of the time.
+
+    Returns:
+        logging.Logger: The root logger.
+    """
+
+    logger = get_logger(name='snnet', log_file=log_file, log_level=log_level)
+
+    return logger
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t", logger=None):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+        self.logger = logger
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        log_msg = [
+            header,
+            '[{0' + space_fmt + '}/{1}]',
+            'eta: {eta}',
+            '{meters}',
+            'time: {time}',
+            'data: {data}'
+        ]
+        if torch.cuda.is_available():
+            log_msg.append('max mem: {memory:.0f}')
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    self.logger.info(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    self.logger.info(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        self.logger.info('{} Total time: {} ({:.4f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+def _load_checkpoint_for_ema(model_ema, checkpoint):
+    """
+    Workaround for ModelEma._load_checkpoint to accept an already-loaded object
+    """
+    mem_file = io.BytesIO()
+    torch.save({'state_dict_ema':checkpoint}, mem_file)
+    mem_file.seek(0)
+    model_ema._load_checkpoint(mem_file)
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print('Not using distributed mode')
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}'.format(
+        args.rank, args.dist_url), flush=True)
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+import json
+def save_on_master_eval_res(log_stats, output_dir):
+    if is_main_process():
+        with open(output_dir, 'a') as f:
+            f.write(json.dumps(log_stats) + "\n")