code of stage1 & 3, remove large files

1_feature_extractor/1_main_training_IB.py

@@ -0,0 +1,624 @@
+# 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 timm.models import create_model
+from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
+from timm.scheduler import create_scheduler
+from timm.optim import create_optimizer
+from timm.utils import NativeScaler, get_state_dict, ModelEma
+from augmentations import collate_data_and_cast_aug
+from datasets import build_dataset
+
+from losses_hint import DistillationLoss
+from samplers import RASampler
+from functools import partial
+
+import importlib
+import utils
+import random
+import math
+from multiprocessing import Value
+from abc import ABC
+
+import sys
+from typing import Iterable, Optional
+from timm.data import Mixup
+from timm.utils import accuracy, ModelEma
+import utils
+import logging
+import torch.distributed as dist
+import os
+
+
+class MaskingGenerator(ABC):
+    def __init__(self, input_size):
+        if not isinstance(input_size, tuple):
+            input_size = (input_size,) * 2
+        self.height, self.width = input_size
+        self.num_patches = self.height * self.width
+
+    def __repr__(self):
+        raise NotImplementedError
+
+    def get_shape(self):
+        return self.height, self.width
+
+    def _mask(self, mask, max_mask_patches):
+        raise NotImplementedError
+
+    def get_none_mask(self):
+        return np.zeros(shape=self.get_shape(), dtype=bool)
+    
+    
+class RandomMaskingGenerator(MaskingGenerator):
+    def __init__(
+        self,
+        input_size,
+    ):
+        """
+        Args:
+            input_size: the size of the token map, e.g., 14x14
+        """
+        super().__init__(input_size)
+
+    def __repr__(self):
+        repr_str = f"Random Generator({self.height}, {self.width})"
+        return repr_str
+
+    def _mask(self, mask, max_mask_patches):
+        return super()._mask(mask, max_mask_patches)
+
+    def __call__(self, num_masking_patches=0):
+        if num_masking_patches <= 0:
+            return np.zeros(shape=self.get_shape(), dtype=bool)
+
+        mask = np.hstack([np.ones(num_masking_patches, dtype=bool),
+                          np.zeros(self.num_patches - num_masking_patches, dtype=bool)])
+        np.random.shuffle(mask)
+        mask = mask.reshape(self.get_shape())
+        return mask
+
+
+def setup_logger(log_dir, rank=0):
+    if rank != 0:
+        return  # 只有主进程（rank 0）配置日志记录器
+    log_formatter = logging.Formatter("%(asctime)s [%(levelname)-5.5s]  %(message)s")
+    root_logger = logging.getLogger()
+    root_logger.setLevel(logging.INFO)
+
+    log_file_handler = logging.FileHandler(log_dir, encoding='utf-8')
+    log_file_handler.setFormatter(log_formatter)
+    root_logger.addHandler(log_file_handler)
+
+    log_stream_handler = logging.StreamHandler(sys.stdout)
+    log_stream_handler.setFormatter(log_formatter)
+    root_logger.addHandler(log_stream_handler)
+
+    logging.info('Logging file is %s' % log_dir)
+
+
+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)
+    parser.add_argument('--target_model', default='deit_base_patch16_224', type=str)
+    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=4e-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
+    
+    # add dataset parameters
+    parser.add_argument('--global_crops_size', '--img_size', default=224, type=int,
+                        help="this should be equal to image size")
+    parser.add_argument('--patch_size', default=16, type=int,
+                        help="patch size for vit patch embedding")
+    
+    # add masking parameter
+    parser.add_argument('--mask_ratio', default=(0.1, 0.5), type=float, nargs='+',
+                        help="mask ratio can be either a value or a range")
+    parser.add_argument('--mask_probability', default=0., type=float,
+                        help="how many samples with be applied with masking")
+    parser.add_argument('--mask_first_n', action='store_true',
+                        help="mask the first n sample to avoid shuffling. Needed for MAE-style encoder")
+    parser.add_argument('--clone_batch', default=1, type=int,
+                        help="how many times to clone the batch for masking (default: 1, not cloning)")
+    
+    # * 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='base', type=str)
+    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="")
+    parser.add_argument('--lambda_token', type=float, default=1.0)
+    parser.add_argument('--lambda_fea', type=float, default=1.0)
+    parser.add_argument('--lambda_patch', type=float, default=1.0)
+    
+    # * Cosub params
+    parser.add_argument('--cosub', action='store_true') 
+    
+    # * Finetuning params
+    parser.add_argument('--finetune', default='', help='finetune from checkpoint')
+    parser.add_argument('--attn-only', action='store_true')
+    parser.add_argument('--weight_inherit', default='') 
+    
+    # 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', 'IMNET_ibot', 'IMNET_ibot_aug', 'IMNET_ibot_fast_aug', 'INAT', 'INAT19', 'IMNET_L', 'IMNET_L_ibot'],
+                        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('--log_dir', default='/data1/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/log', 
+                        type=str, help='saving logging info every 20 iters')
+    parser.add_argument('--device', default='cuda',
+                        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('--distributed', action='store_true', default=False, help='Enabling distributed training')
+    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')
+    return parser
+
+
+def show_learnable_params(model):
+    enabled = set()
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            enabled.add(name)
+    # print("Parameters to be updated: ")
+    logging.info("Parameters to be updated: ")
+    for each in enabled:
+        # print('\t{}\n'.format(str(each)))
+        logging.info('\t{}\n'.format(str(each)))
+    # print('\n')
+    logging.info('\n')
+
+def show_unlearnable_params(model):
+    disabled = set()
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            disabled.add(name)
+    
+    logging.info("Parameters that are not being updated: ")
+    for each in disabled:
+        logging.info('\t{}'.format(str(each)))
+    logging.info('\n')
+
+def main(args):
+    utils.init_distributed_mode(args)
+
+    print(args)
+
+    device = torch.device(args.device)
+     # 获取当前进程的 rank
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    # set up logger
+    os.makedirs(args.log_dir, exist_ok=True)
+    setup_logger(args.log_dir + '/' + time.strftime('%Y%m%d_%H%M%S') + '.log', rank)
+    logging.info('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
+    logging.info("{}".format(args).replace(', ', ',\n') + '\n')
+
+    # 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
+
+    dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
+    logging.info(dataset_train)
+
+    if args.distributed:
+        num_tasks = utils.get_world_size()
+        global_rank = utils.get_rank()
+        if args.repeated_aug:
+            sampler_train = RASampler(
+                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+            )
+        else:
+            sampler_train = torch.utils.data.DistributedSampler(
+                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+            )
+    else:
+        sampler_train = torch.utils.data.RandomSampler(dataset_train)
+    logging.info("Sampler_train = %s" % str(sampler_train))
+        
+    n_tokens = (args.global_crops_size // args.patch_size) ** 2
+    mask_generator = RandomMaskingGenerator(
+        input_size=args.global_crops_size // args.patch_size,
+    )
+
+    collate_fn = partial(
+        collate_data_and_cast_aug,
+        mask_ratio=args.mask_ratio,
+        mask_probability=args.mask_probability,
+        dtype=torch.half,   # half precision
+        n_tokens=n_tokens,
+        mask_first_n=args.mask_first_n,
+        mask_generator=mask_generator,
+        clone_batch=args.clone_batch,
+    )
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train, sampler=sampler_train,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+        collate_fn=collate_fn,
+    )
+
+    mixup_fn = None
+
+    print(f"Creating model: {args.model}")  # models_proteus_dinov2
+    meta_arch_module = importlib.import_module(args.model)
+    MetaArch = meta_arch_module.MetaArch
+
+    model = MetaArch(args)
+    logging.info("Model = %s" % str(model))
+
+    if args.finetune:
+        checkpoint = torch.load(args.finetune, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            pretrained_dict = checkpoint['state_dict']
+        elif 'model' in checkpoint:
+            pretrained_dict = checkpoint['model']
+        else:
+            pretrained_dict = checkpoint
+
+        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
+        # print('missing_keys: ', missing_keys)
+        # print('unexpected_keys: ', unexpected_keys)
+        logging.info('Finetuning from %s' % args.finetune)
+        logging.info('missing_keys: %s' % str(missing_keys))
+        logging.info('unexpected_keys: %s' % str(unexpected_keys))
+        
+    if args.attn_only:
+        for name_p,p in model.named_parameters():
+            if '.attn.' in name_p:
+                p.requires_grad = True
+            else:
+                p.requires_grad = False
+        try:
+            model.head.weight.requires_grad = True
+            model.head.bias.requires_grad = True
+        except:
+            model.fc.weight.requires_grad = True
+            model.fc.bias.requires_grad = True
+        try:
+            model.pos_embed.requires_grad = True
+        except:
+            print('no position encoding')
+        try:
+            for p in model.patch_embed.parameters():
+                p.requires_grad = False
+        except:
+            print('no patch embed')
+            
+    model.to(device)
+
+    model_ema = None
+    if args.model_ema:
+        # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
+        model_ema = ModelEma(
+            model.student.backbone,
+            decay=args.model_ema_decay,
+            device='cpu' if args.model_ema_force_cpu else '',
+            resume='')
+
+    model_without_ddp = model
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
+        model_without_ddp = model.module
+    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    # print('number of params:', n_parameters)
+    logging.info('number of params: %s' % n_parameters)
+    
+    if not args.unscale_lr:
+        logging.info('base lr = %s' % args.lr)
+        linear_scaled_lr = args.lr * args.batch_size * utils.get_world_size() / 512.0
+        args.lr = linear_scaled_lr
+        logging.info('actural lr = %s' % linear_scaled_lr)
+    
+    optimizer = create_optimizer(args, model_without_ddp)
+    loss_scaler = NativeScaler()
+
+    lr_scheduler, _ = create_scheduler(args, optimizer)
+
+    output_dir = Path(args.output_dir)
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume, map_location='cpu')
+
+        model_without_ddp.load_state_dict(checkpoint['model'])
+        if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if args.model_ema:
+                utils._load_checkpoint_for_ema(model_ema, checkpoint['model_ema'])
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+        lr_scheduler.step(args.start_epoch)
+        logging.info('Resuming from %s' % args.resume)
+
+    # print(f"Start training for {args.epochs} epochs")
+    logging.info("Start training for %s epochs" % args.epochs)
+    start_time = time.time()
+    max_accuracy = 0.0
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            data_loader_train.sampler.set_epoch(epoch)
+        
+        if epoch < 5:
+            # 前5个epoch仅放开 entropy model 的参数
+            for name, param in model.named_parameters():
+                if 'info_bottleneck' in name:
+                    param.requires_grad = True
+                else:
+                    param.requires_grad = False
+            if epoch == 0:
+                show_learnable_params(model)
+        else:
+             # 其余epoch放开所有参数，但固定model.teacher的参数
+            for name, param in model.named_parameters():
+                    param.requires_grad = True
+            for name, param in model.named_parameters():
+                if 'teacher' in name:
+                    param.requires_grad = False
+            if epoch == 5:
+                show_unlearnable_params(model)
+        
+        train_stats = train_one_epoch(
+            model, data_loader_train,
+            optimizer, device, epoch, loss_scaler,
+            args.clip_grad, model_ema, mixup_fn,
+            set_training_mode=args.train_mode,  # keep in eval mode for deit finetuning / train mode for training and deit III finetuning
+            args = args,
+        )
+
+        lr_scheduler.step(epoch)
+        if args.output_dir:
+            checkpoint_paths = [output_dir / 'checkpoint.pth']
+            for checkpoint_path in checkpoint_paths:
+                utils.save_on_master({
+                    'model': model_without_ddp.state_dict(),
+                    'optimizer': optimizer.state_dict(),
+                    'lr_scheduler': lr_scheduler.state_dict(),
+                    'epoch': epoch,
+                    'model_ema': get_state_dict(model_ema),
+                    'scaler': loss_scaler.state_dict(),
+                    'args': args,
+                }, checkpoint_path)
+        if args.output_dir and (epoch % 5 == 0 or epoch + 1 == args.epochs):
+            checkpoint_path = output_dir / f'checkpoint{epoch:04}.pth'
+            utils.save_on_master({
+                'model': model_without_ddp.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'lr_scheduler': lr_scheduler.state_dict(),
+                'epoch': epoch,
+                'model_ema': get_state_dict(model_ema),
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }, checkpoint_path)
+
+        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+                     'epoch': epoch,
+                     'n_parameters': n_parameters}
+        
+        if args.output_dir and utils.is_main_process():
+            with (output_dir / "log.txt").open("a") as f:
+                f.write(json.dumps(log_stats) + "\n")
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    # print('Training time {}'.format(total_time_str))
+    logging.info('Training time %s' % total_time_str)
+
+
+def train_one_epoch(model: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
+                    model_ema: Optional[ModelEma] = None, mixup_fn: Optional[Mixup] = None,
+                    set_training_mode=True, args = None):
+    model.train(set_training_mode)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 20
+    
+    loader_len = len(data_loader)
+
+    for data_iter_step, inputs_dict in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+
+        for k, v in inputs_dict.items():
+            if isinstance(v, torch.Tensor):
+                inputs_dict[k] = v.to(device, non_blocking=True)
+        
+        with torch.cuda.amp.autocast():
+            loss_dict = model(inputs_dict)
+        
+        loss = loss_dict["loss"]
+        patch_loss = loss_dict["patch_loss"]
+        fea_loss = loss_dict["fea_loss"]
+        token_loss = loss_dict["token_loss"]
+        bpp_loss = loss_dict["bpp_loss"]
+        task_loss = loss_dict["task_loss"]
+
+        patch_loss_value = patch_loss.item()
+        token_loss_value = token_loss.item()
+        fea_loss_value = fea_loss.item()
+        bpp_loss_value = bpp_loss.item()
+        task_loss_value = task_loss.item()
+        loss_value = loss.item()
+        
+
+        if not math.isfinite(loss_value):
+            # print("Loss is {}, stopping training".format(loss_value))
+            logging.info("Loss is %s, stopping training" % loss_value)
+            logging.info("bpp_loss is {}, patch_loss is {}, token_loss is {}, fea_loss is {}".format(bpp_loss_value, patch_loss_value, token_loss_value, fea_loss_value))
+            sys.exit(1)
+
+        optimizer.zero_grad()
+
+        # this attribute is added by timm on one optimizer (adahessian)
+        is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
+        loss_scaler(loss, optimizer, clip_grad=max_norm,
+                    parameters=model.parameters(), create_graph=is_second_order)
+
+        torch.cuda.synchronize()
+        if model_ema is not None:
+            model_ema.update(model.module.student.backbone)
+
+        metric_logger.update(loss=loss_value)
+        metric_logger.update(task_loss=task_loss_value)
+        metric_logger.update(bpp_loss=bpp_loss_value)
+        metric_logger.update(patch_loss=patch_loss_value)
+        metric_logger.update(token_loss=token_loss_value)
+        metric_logger.update(fea_loss=fea_loss_value)
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    # print("Averaged stats:", metric_logger)
+    logging.info("Averaged stats: {}".format(metric_logger))
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('DeiT training and evaluation script', parents=[get_args_parser()])
+    args = parser.parse_args()
+    if args.output_dir:
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    main(args)

1_feature_extractor/1_training_IB.sh

@@ -0,0 +1,40 @@
+
+
+#### access DINOv2
+export CUDA_VISIBLE_DEVICES=0,1,2;
+
+python -m torch.distributed.launch --nproc_per_node=3 --use_env 1_main_training_IB.py \
+    --batch-size 48 --warmup-epochs 5 --epochs 200 \
+    --data-set IMNET --data-path '/data1/datasets/imagenet_fold' \
+    --teacher-model vit_large --target_model vit_base --model models_proteus_dinov2 \
+    --patch_size 14 --mask_probability 0.5 --mask_ratio 0.5 --mask_first_n \
+    --lambda_token 1.0 --lambda_fea 1.05 --lambda_patch 1.05 \
+    --resume "/data0/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/checkpoint0160.pth" \
+    --log_dir '/data0/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/log/' \
+    --output_dir log/DINOv2_training;
+
+
+
+#### access SynCLR
+
+# python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py \
+#     --batch-size 128 --warmup-epochs 5 --epochs 300 \
+#     --data-set IMNET --data-path imagenet_path \
+#     --teacher-model vit_large --target_model vit_base --model models_proteus_synclr \
+#     --teacher-path pretrained_synclr_path \
+#     --patch_size 14 --mask_probability 0.5 --mask_ratio 0.5 --mask_first_n \
+#     --lambda_token 1.0 --lambda_fea 1.0 --lambda_patch 1.0 \
+#     --output_dir log/SynCLR_training;
+
+
+
+#### access CLIP
+
+# python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py \
+#     --batch-size 128 --warmup-epochs 5 --epochs 300 \
+#     --data-set IMNET --data-path imagenet_path \
+#     --teacher-model vit_large --target_model vit_base --model models_proteus_clip \
+#     --teacher-path pretrained_clip_path \
+#     --patch_size 14 --mask_probability 0.5 --mask_ratio 0.5 --mask_first_n \
+#     --lambda_token 1.0 --lambda_fea 0.0 --lambda_patch 0.0 \
+#     --output_dir log/CLIP_training;

1_feature_extractor/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Yunpeng Qi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

1_feature_extractor/README copy.md

@@ -0,0 +1,24 @@
+# Pre-training on ImageNet-1K
+
+
+## Installation
+Please follow the installation instructions in [DINOv2](https://github.com/facebookresearch/dinov2/tree/main?tab=readme-ov-file#installation) and install timm==0.9.16 as well.
+
+## Dataset
+We prepare ImageNet-1K following the instructions in [DeiT](https://github.com/facebookresearch/deit/blob/main/README_deit.md#data-preparation).
+
+## Training
+1. Specify the directory of datasets with `data-path` in the training script `run_pretrain.sh`.
+2. Use the `teacher-model` and `target_model` parameters to select the appropriate teacher and student models.
+3. Specify the model choices with `model` to choose from DINOv2, SynCLR, CLIP.
+4. For SynCLR and CLIP training, use the `teacher-path` parameter to indicate the path to the pre-trained teacher model.
+5. Simply run the training script as follows:
+
+   ```
+   bash run_pretrain.sh
+   ```
+
+
+## Acknowledgment
+
+This part is heavily build upon [DeiT](https://github.com/facebookresearch/deit?tab=readme-ov-file), [DINOv2](https://github.com/facebookresearch/dinov2), [SynCLR](https://github.com/google-research/syn-rep-learn/tree/main/SynCLR). We gratefully thank the authors for their wonderful works.

1_feature_extractor/README.md

@@ -0,0 +1,17 @@
+# 1_feature_extractor
+training information bottleneck
+
+IF Model: https://huggingface.co/Qiyp/1_feature_extractor
+
+Installation:
+Clone the repository and then use the provided requirements.txt to install the dependencies:
+pip install -r requirements.txt
+
+put the proteus_vitb_backbone.pth into ./ckpt
+
+load pretrained Feature Extractor for further training: --resume + Feature Extractor path
+
+load pretrained Feature Extractor for feature extractraction: --finetune + Feature Extractor path
+
+training scripts: 1_training_IB.sh
+vis feature's semantic information: train_dec.sh

1_feature_extractor/augmentations.py

@@ -0,0 +1,94 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import random
+from torchvision import transforms
+import torch
+
+logger = logging.getLogger("dinov2")
+
+
+
+def collate_data_and_cast_aug(
+    samples_list,
+    mask_ratio,
+    mask_probability,
+    dtype,
+    n_tokens=None,
+    mask_first_n=False,
+    mask_generator=None,
+    clone_batch=1,
+):
+    # dtype = torch.half  # TODO: Remove
+
+    n_global_crops = 1
+
+    assert n_global_crops > 0, "global crops number should be > 0"
+    collated_global_crops = torch.stack([s[i] for i in range(n_global_crops) for s in samples_list])
+
+    labels = [s[1] for s in samples_list]
+    labels = torch.LongTensor(labels)
+    collated_global_labels = labels.repeat(n_global_crops)
+
+    B = len(collated_global_crops)
+    N = n_tokens
+    n_samples_masked = int(B * mask_probability)
+
+    masks_list = []
+    upperbound = 0
+
+    masks_enc = torch.full((1,), 0, dtype=torch.int32)
+    masks_pred = torch.full((1,), 0, dtype=torch.int32)
+    # specify the number of masks to append
+    number_masks = n_samples_masked * clone_batch
+    # do per-sample masking
+    if isinstance(mask_ratio, (tuple, list)) and len(mask_ratio) == 2:
+        probs = torch.linspace(*mask_ratio, number_masks + 1)
+        for i in range(0, number_masks):
+            prob_min = probs[i]
+            prob_max = probs[i + 1]
+            masks_list.append(torch.BoolTensor(mask_generator(int(N * random.uniform(prob_min, prob_max)))))
+            upperbound += int(N * prob_max)
+    else:
+        mask_ratio = mask_ratio[0]
+        # apply the same mask ratio to all images
+        for i in range(0, number_masks):
+            masks_list.append(torch.BoolTensor(mask_generator(int(N * mask_ratio))))
+            upperbound += int(N * mask_ratio)
+
+    # append masks for unmasked samples
+    for i in range(n_samples_masked, B):
+        # masks_list.append(torch.BoolTensor(mask_generator(0)))
+        masks_list.append(torch.BoolTensor(mask_generator.get_none_mask()))
+
+    if not mask_first_n and mask_probability > 0.0:  # shuffle masking -- not shuffling for mae-style
+        random.shuffle(masks_list)
+
+    collated_masks = torch.stack(masks_list).flatten(1)
+    mask_indices_list = collated_masks.flatten().nonzero().flatten()
+
+    masks_weight = (1 / collated_masks.sum(-1).clamp(min=1.0)).unsqueeze(-1).expand_as(collated_masks)[collated_masks]
+
+    return {
+        "collated_global_crops": collated_global_crops.to(dtype),
+        "collated_global_labels": collated_global_labels,
+        "collated_masks": collated_masks,
+        "mask_indices_list": mask_indices_list,
+        "masks_weight": masks_weight,
+        "upperbound": upperbound,
+        "n_masked_patches": torch.full((1,), fill_value=mask_indices_list.shape[0], dtype=torch.long),
+        "masks_enc": masks_enc,
+        "masks_pred": masks_pred,
+    }

1_feature_extractor/datasets.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+import os
+import json
+from torchvision.datasets import DatasetFolder
+from torchvision.io import read_image
+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
+from PIL import Image
+
+
+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)

1_feature_extractor/fast_vis.sh

@@ -0,0 +1,37 @@
+
+
+#### access DINOv2
+export CUDA_VISIBLE_DEVICES=6;
+
+python fast_vis_settings_all.py \
+    --batch-size 64 --warmup-epochs 5 --epochs 300 \
+    --data-set IMNET --data-path '/data1/datasets/imagenet_fold' \
+    --teacher-model vit_large --target_model vit_base --model models_proteus_dinov2 \
+    --patch_size 14 --mask_probability 0.5 --mask_ratio 0.5 --mask_first_n \
+    --lambda_token 1.0 --lambda_fea 1.0 --lambda_patch 1.0 \
+    --finetune "/data0/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/checkpoint0160.pth" \
+    --output_dir log/DINOv2_training;
+
+#### access SynCLR
+
+# python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py \
+#     --batch-size 128 --warmup-epochs 5 --epochs 300 \
+#     --data-set IMNET --data-path imagenet_path \
+#     --teacher-model vit_large --target_model vit_base --model models_proteus_synclr \
+#     --teacher-path pretrained_synclr_path \
+#     --patch_size 14 --mask_probability 0.5 --mask_ratio 0.5 --mask_first_n \
+#     --lambda_token 1.0 --lambda_fea 1.0 --lambda_patch 1.0 \
+#     --output_dir log/SynCLR_training;
+
+
+
+#### access CLIP
+
+# python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py \
+#     --batch-size 128 --warmup-epochs 5 --epochs 300 \
+#     --data-set IMNET --data-path imagenet_path \
+#     --teacher-model vit_large --target_model vit_base --model models_proteus_clip \
+#     --teacher-path pretrained_clip_path \
+#     --patch_size 14 --mask_probability 0.5 --mask_ratio 0.5 --mask_first_n \
+#     --lambda_token 1.0 --lambda_fea 0.0 --lambda_patch 0.0 \
+#     --output_dir log/CLIP_training;

1_feature_extractor/fast_vis_proteus_feats.py

@@ -0,0 +1,98 @@
+import math
+import matplotlib.pyplot as plt
+import numpy as np
+import einops
+import torch
+from PIL import Image
+import torchvision
+from PIL import Image
+import models_dinov2
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+patch_size = 14
+
+# feat_extractor = getattr(models_dinov2, 'vit_base')
+feat_extractor = getattr(models_dinov2, 'vit_large')
+model = feat_extractor(img_size=224,
+            patch_size=14,
+            init_values=1.0,
+            ffn_layer='mlp',
+            block_chunks=0,
+            num_register_tokens=0,
+            interpolate_antialias=False,
+            interpolate_offset=0.1)
+
+# checkpoint_path = '/data0/qiyp/Proteus-pytorch/pretrain/ckpt/proteus_vitb_backbone.pth'  # 替换为实际的检查点路径
+checkpoint_path = '/data0/qiyp/Proteus-pytorch/pretrain/ckpt/proteus_vitl_backbone.pth'  # 替换为实际的检查点路径
+# 加载检查点
+checkpoint = torch.load(checkpoint_path, map_location='cpu')
+
+# 加载模型权重
+if 'state_dict' in checkpoint:
+    pretrained_dict = checkpoint['state_dict']
+elif 'model' in checkpoint:
+    pretrained_dict = checkpoint['model']
+else:
+    pretrained_dict = checkpoint
+
+# 只加载与学生模型相关的部分
+model.load_state_dict(pretrained_dict, strict=False)
+model.to(device)
+patch_h = 224 // 14
+patch_w = 224 // 14
+feat_dim = 768
+
+def visualize_features(features, output_path='./feature_visualization.png'):
+    # Assuming features are of shape (batch_size, num_features, height, width)
+    batch_size, num_features, height, width = features.shape
+    
+    # Normalize the feature maps to the range [0, 1]
+    vis = features.mean(dim=1, keepdim=True)
+    vis = vis - vis.min()
+    vis = vis / vis.max()
+    
+    # Squeeze the channel dimension
+    vis = vis.squeeze(1).cpu().detach().numpy()
+    
+    # Apply a colormap (e.g., viridis) to convert it to RGB
+    vis_colored = np.zeros((batch_size, height, width, 3))
+    for i in range(batch_size):
+        vis_colored[i] = plt.cm.viridis(vis[i])[:, :, :3]  # Drop the alpha channel
+    
+    # Convert vis_colored to a tensor and save using torchvision
+    vis_colored = torch.tensor(vis_colored).permute(0, 3, 1, 2)  # Convert to (batch, channels, height, width)
+    
+    # Save the image
+    torchvision.utils.save_image(vis_colored, output_path, normalize=True)
+
+from torchvision import transforms
+
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),  # 调整图像大小
+    transforms.ToTensor(),  # 转换为tensor
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
+])
+
+images = [
+    Image.open("/data0/qiyp/mae/imagenet-1k-samples/0-anime_boy_sticker__holding_kitten__happy.png"),
+    Image.open("/data0/qiyp/mae/imagenet-1k-samples/62-Deadpool_minion.png"),
+    Image.open("/data0/qiyp/mae/imagenet-1k-samples/79-with_Wooden_carved_bear__salmon_and_gold_mini_ball_surround_the_blank_signboard__illustrate.png"),
+    Image.open("/data0/qiyp/mae/imagenet-1k-samples/99-Akira_toriyama_motorbike__cheatah__puma__japanese_classic_car__collectable_figure__shiny_plastic_.png"),
+    Image.open("/data0/qiyp/mae/imagenet-1k-samples/124-crowded_1920s_Chicago_street_with_lots_of_model_T_cars_and_people.png"),
+    Image.open("/data0/qiyp/mae/imagenet-1k-samples/157-steampunk_girl_with_pink_hair_riding_in_a_hot_air_balloon__hot_air_balloon_resembles_gold_and_si.png"),
+    Image.open("/data0/qiyp/mae/imagenet-1k-samples/ILSVRC2012_val_00008636.png"),
+    Image.open("/data0/qiyp/mae/imagenet-1k-samples/ILSVRC2012_val_00010240.png"),
+          ]
+# inputs = processor(images=images, return_tensors="pt", padding=True).to(device)
+tensors = [transform(img) for img in images]
+batched_tensors = torch.stack(tensors).to(device)
+with torch.no_grad():
+    outputs = model(batched_tensors, is_training=True)
+    features = outputs['x_norm_patchtokens']  # (batch_size, num_patches, feat_dim)
+    print(features.shape)
+
+    features = features.view(-1, patch_h, patch_w, features.shape[2])  # [B, h, w, c]
+    features = features.permute(0, 3, 1, 2)
+    visualize_features(features)
+
+    # pooled_output = outputs.pooler_output  # pooled CLS states. 

1_feature_extractor/fast_vis_settings_all.py

@@ -0,0 +1,548 @@
+# 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 timm.models import create_model
+from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
+from timm.scheduler import create_scheduler
+from timm.optim import create_optimizer
+from timm.utils import NativeScaler, get_state_dict, ModelEma
+from augmentations import collate_data_and_cast_aug
+from datasets import build_dataset
+
+from losses_hint import DistillationLoss
+from samplers import RASampler
+from functools import partial
+
+import importlib
+import utils
+import random
+import math
+from multiprocessing import Value
+from abc import ABC
+
+import sys
+from typing import Iterable, Optional
+from timm.data import Mixup
+from timm.utils import accuracy, ModelEma
+import utils
+import os
+# os.environ["CUDA_VISIBLE_DEVICES"] = "5"
+
+
+class MaskingGenerator(ABC):
+    def __init__(self, input_size):
+        if not isinstance(input_size, tuple):
+            input_size = (input_size,) * 2
+        self.height, self.width = input_size
+        self.num_patches = self.height * self.width
+
+    def __repr__(self):
+        raise NotImplementedError
+
+    def get_shape(self):
+        return self.height, self.width
+
+    def _mask(self, mask, max_mask_patches):
+        raise NotImplementedError
+
+    def get_none_mask(self):
+        return np.zeros(shape=self.get_shape(), dtype=bool)
+    
+    
+    
+class RandomMaskingGenerator(MaskingGenerator):
+    def __init__(
+        self,
+        input_size,
+    ):
+        """
+        Args:
+            input_size: the size of the token map, e.g., 14x14
+        """
+        super().__init__(input_size)
+
+    def __repr__(self):
+        repr_str = f"Random Generator({self.height}, {self.width})"
+        return repr_str
+
+    def _mask(self, mask, max_mask_patches):
+        return super()._mask(mask, max_mask_patches)
+
+    def __call__(self, num_masking_patches=0):
+        if num_masking_patches <= 0:
+            return np.zeros(shape=self.get_shape(), dtype=bool)
+
+        mask = np.hstack([np.ones(num_masking_patches, dtype=bool),
+                          np.zeros(self.num_patches - num_masking_patches, dtype=bool)])
+        np.random.shuffle(mask)
+        mask = mask.reshape(self.get_shape())
+        return mask
+
+
+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)
+    parser.add_argument('--target_model', default='deit_base_patch16_224', type=str)
+    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
+    
+    # add dataset parameters
+    parser.add_argument('--global_crops_size', '--img_size', default=224, type=int,
+                        help="this should be equal to image size")
+    parser.add_argument('--patch_size', default=16, type=int,
+                        help="patch size for vit patch embedding")
+    
+    # add masking parameter
+    parser.add_argument('--mask_ratio', default=(0.1, 0.5), type=float, nargs='+',
+                        help="mask ratio can be either a value or a range")
+    parser.add_argument('--mask_probability', default=0., type=float,
+                        help="how many samples with be applied with masking")
+    parser.add_argument('--mask_first_n', action='store_true',
+                        help="mask the first n sample to avoid shuffling. Needed for MAE-style encoder")
+    parser.add_argument('--clone_batch', default=1, type=int,
+                        help="how many times to clone the batch for masking (default: 1, not cloning)")
+    
+    # * 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='base', type=str)
+    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="")
+    parser.add_argument('--lambda_token', type=float, default=1.0)
+    parser.add_argument('--lambda_fea', type=float, default=1.0)
+    parser.add_argument('--lambda_patch', type=float, default=1.0)
+    
+    # * Cosub params
+    parser.add_argument('--cosub', action='store_true') 
+    
+    # * Finetuning params
+    parser.add_argument('--finetune', default='', help='finetune from checkpoint')
+    parser.add_argument('--attn-only', action='store_true')
+    parser.add_argument('--weight_inherit', default='') 
+    
+    # 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', 'IMNET_ibot', 'IMNET_ibot_aug', 'IMNET_ibot_fast_aug', 'INAT', 'INAT19', 'IMNET_L', 'IMNET_L_ibot'],
+                        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='cuda',
+                        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('--distributed', action='store_true', default=False, help='Enabling distributed training')
+    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')
+    return parser
+
+import torchvision
+import matplotlib.pyplot as plt
+import torchvision.transforms as transforms
+def visualize_features(features, output_path='./feature_visualization_w_ib.png'):
+    # Assuming features are of shape (batch_size, num_features, height, width)
+    batch_size, num_features, height, width = features.shape
+    
+    # Normalize the feature maps to the range [0, 1]
+    vis = features.mean(dim=1, keepdim=True)
+    vis = vis - vis.min()
+    vis = vis / vis.max()
+    
+    # Squeeze the channel dimension
+    vis = vis.squeeze(1).cpu().detach().numpy()
+    
+    # Apply a colormap (e.g., viridis) to convert it to RGB
+    vis_colored = np.zeros((batch_size, height, width, 3))
+    for i in range(batch_size):
+        vis_colored[i] = plt.cm.viridis(vis[i])[:, :, :3]  # Drop the alpha channel
+    
+    # Convert vis_colored to a tensor and save using torchvision
+    vis_colored = torch.tensor(vis_colored).permute(0, 3, 1, 2)  # Convert to (batch, channels, height, width)
+    
+    # Save the image
+    torchvision.utils.save_image(vis_colored, output_path, normalize=True)
+
+def save_original_images(tensors, output_path='./original_images.png'):
+    # 将归一化反转
+    unnormalize = transforms.Normalize(
+        mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225],
+        std=[1/0.229, 1/0.224, 1/0.225]
+    )
+    unnormalized_tensors = [unnormalize(tensor) for tensor in tensors]
+    unnormalized_batch = torch.stack(unnormalized_tensors)
+    torchvision.utils.save_image(unnormalized_batch, output_path, nrow=4, normalize=True)
+
+def main(args):
+    utils.init_distributed_mode(args)
+
+    print(args)
+
+    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
+
+    dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
+
+    if args.distributed:
+        num_tasks = utils.get_world_size()
+        global_rank = utils.get_rank()
+        if args.repeated_aug:
+            sampler_train = RASampler(
+                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+            )
+        else:
+            sampler_train = torch.utils.data.DistributedSampler(
+                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+            )
+    else:
+        sampler_train = torch.utils.data.RandomSampler(dataset_train)
+        
+    n_tokens = (args.global_crops_size // args.patch_size) ** 2
+    mask_generator = RandomMaskingGenerator(
+        input_size=args.global_crops_size // args.patch_size,
+    )
+
+    collate_fn = partial(
+        collate_data_and_cast_aug,
+        mask_ratio=args.mask_ratio,
+        mask_probability=args.mask_probability,
+        dtype=torch.half,   # half precision
+        n_tokens=n_tokens,
+        mask_first_n=args.mask_first_n,
+        mask_generator=mask_generator,
+        clone_batch=args.clone_batch,
+    )
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train, sampler=sampler_train,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+        collate_fn=collate_fn,
+    )
+
+    mixup_fn = None
+
+    print(f"Creating model: {args.model}")
+    meta_arch_module = importlib.import_module(args.model)
+    MetaArch = meta_arch_module.MetaArch
+
+    model = MetaArch(args)
+
+    if args.finetune:
+        checkpoint = torch.load(args.finetune, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            pretrained_dict = checkpoint['state_dict']
+        elif 'model' in checkpoint:
+            pretrained_dict = checkpoint['model']
+        else:
+            pretrained_dict = checkpoint
+
+        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
+        print('missing_keys: ', missing_keys)
+        print('unexpected_keys: ', unexpected_keys)
+        
+    if args.attn_only:
+        for name_p,p in model.named_parameters():
+            if '.attn.' in name_p:
+                p.requires_grad = True
+            else:
+                p.requires_grad = False
+        try:
+            model.head.weight.requires_grad = True
+            model.head.bias.requires_grad = True
+        except:
+            model.fc.weight.requires_grad = True
+            model.fc.bias.requires_grad = True
+        try:
+            model.pos_embed.requires_grad = True
+        except:
+            print('no position encoding')
+        try:
+            for p in model.patch_embed.parameters():
+                p.requires_grad = False
+        except:
+            print('no patch embed')
+            
+    model.to(device)
+
+    model_ema = None
+    if args.model_ema:
+        # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
+        model_ema = ModelEma(
+            model.student.backbone,
+            decay=args.model_ema_decay,
+            device='cpu' if args.model_ema_force_cpu else '',
+            resume='')
+
+    model_without_ddp = model
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
+        model_without_ddp = model.module
+    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print('number of params:', n_parameters)
+    
+    if not args.unscale_lr:
+        linear_scaled_lr = args.lr * args.batch_size * utils.get_world_size() / 512.0
+        args.lr = linear_scaled_lr
+    
+    optimizer = create_optimizer(args, model_without_ddp)
+    loss_scaler = NativeScaler()
+
+    lr_scheduler, _ = create_scheduler(args, optimizer)
+
+    output_dir = Path(args.output_dir)
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume, map_location='cpu')
+
+        model_without_ddp.load_state_dict(checkpoint['model'])
+        if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if args.model_ema:
+                utils._load_checkpoint_for_ema(model_ema, checkpoint['model_ema'])
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+        lr_scheduler.step(args.start_epoch)
+    
+
+    from torchvision import transforms
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),  # 调整图像大小
+        transforms.ToTensor(),  # 转换为tensor
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
+    ])
+    from PIL import Image
+    images = [
+        Image.open("/data0/qiyp/mae/imagenet-1k-samples/0-anime_boy_sticker__holding_kitten__happy.png"),
+        Image.open("/data0/qiyp/mae/imagenet-1k-samples/62-Deadpool_minion.png"),
+        Image.open("/data0/qiyp/mae/imagenet-1k-samples/79-with_Wooden_carved_bear__salmon_and_gold_mini_ball_surround_the_blank_signboard__illustrate.png"),
+        Image.open("/data0/qiyp/mae/imagenet-1k-samples/99-Akira_toriyama_motorbike__cheatah__puma__japanese_classic_car__collectable_figure__shiny_plastic_.png"),
+        Image.open("/data0/qiyp/mae/imagenet-1k-samples/124-crowded_1920s_Chicago_street_with_lots_of_model_T_cars_and_people.png"),
+        Image.open("/data0/qiyp/mae/imagenet-1k-samples/157-steampunk_girl_with_pink_hair_riding_in_a_hot_air_balloon__hot_air_balloon_resembles_gold_and_si.png"),
+        Image.open("/data0/qiyp/mae/imagenet-1k-samples/ILSVRC2012_val_00008636.png"),
+        Image.open("/data0/qiyp/mae/imagenet-1k-samples/ILSVRC2012_val_00010240.png"),
+        ]
+
+    tensors = [transform(img) for img in images]
+    batched_tensors = torch.stack(tensors).to(device)
+    save_original_images(batched_tensors, output_path='./original_images.png')
+    with torch.no_grad():
+        outputs = model.student.backbone(batched_tensors, is_training=True)
+        features = outputs['x_norm_patchtokens']  # (batch_size, num_patches, feat_dim)
+        print(features.shape)
+        features, _ = model.info_bottleneck(features, is_training=False)
+
+        features = features.view(-1, 16, 16, features.shape[2])  # [B, h, w, c]
+        features = features.permute(0, 3, 1, 2)
+        visualize_features(features)
+
+
+def train_one_epoch(model: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
+                    model_ema: Optional[ModelEma] = None, mixup_fn: Optional[Mixup] = None,
+                    set_training_mode=True, args = None):
+    model.train(set_training_mode)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 10
+    
+    loader_len = len(data_loader)
+
+    for data_iter_step, inputs_dict in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+
+        for k, v in inputs_dict.items():
+            if isinstance(v, torch.Tensor):
+                inputs_dict[k] = v.to(device, non_blocking=True)
+        
+        with torch.cuda.amp.autocast():
+            loss_dict = model(inputs_dict)
+        
+        loss = loss_dict["loss"]
+        patch_loss = loss_dict["patch_loss"]
+        fea_loss = loss_dict["fea_loss"]
+        token_loss = loss_dict["token_loss"]
+
+        patch_loss_value = patch_loss.item()
+        token_loss_value = token_loss.item()
+        fea_loss_value = fea_loss.item()
+        loss_value = loss.item()
+        
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        optimizer.zero_grad()
+
+        # this attribute is added by timm on one optimizer (adahessian)
+        is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
+        loss_scaler(loss, optimizer, clip_grad=max_norm,
+                    parameters=model.parameters(), create_graph=is_second_order)
+
+        torch.cuda.synchronize()
+        if model_ema is not None:
+            model_ema.update(model.module.student.backbone)
+
+        metric_logger.update(loss=loss_value)
+        metric_logger.update(patch_loss=patch_loss_value)
+        metric_logger.update(token_loss=token_loss_value)
+        metric_logger.update(fea_loss=fea_loss_value)
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('DeiT training and evaluation script', parents=[get_args_parser()])
+    args = parser.parse_args()
+    if args.output_dir:
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    main(args)

1_feature_extractor/log/DINOv2_training/log.txt

@@ -0,0 +1,203 @@
+{"train_lr": 1.0000000000000353e-06, "train_loss": 9.738612308347825, "train_task_loss": 4.1162851987411075, "train_bpp_loss": 28.111635084060744, "train_patch_loss": 2.1926947419615765, "train_token_loss": 0.6686155230319091, "train_fea_loss": 1.254974935296104, "epoch": 0, "n_parameters": 144845568}
+{"train_lr": 1.0000000000000353e-06, "train_loss": 9.232701392577802, "train_task_loss": 3.6589209317660734, "train_bpp_loss": 27.86890183459464, "train_patch_loss": 1.9597956623069102, "train_token_loss": 0.6684419002523906, "train_fea_loss": 1.030683369156149, "epoch": 1, "n_parameters": 144845568}
+{"train_lr": 7.579999999999412e-05, "train_loss": 7.21169285546485, "train_task_loss": 3.139855162607585, "train_bpp_loss": 20.35918810360914, "train_patch_loss": 1.6089605613649607, "train_token_loss": 0.6681283143200843, "train_fea_loss": 0.8627662869969861, "epoch": 2, "n_parameters": 144845568}
+{"train_lr": 0.0001506000000000007, "train_loss": 4.554179431234809, "train_task_loss": 3.242361048226067, "train_bpp_loss": 6.559091770915772, "train_patch_loss": 1.606790432570983, "train_token_loss": 0.6681644691438745, "train_fea_loss": 0.9674061452158444, "epoch": 3, "n_parameters": 144845568}
+{"train_lr": 0.00022540000000000762, "train_loss": 3.735120667863807, "train_task_loss": 3.387666613328085, "train_bpp_loss": 1.7372702604485788, "train_patch_loss": 1.6483151540636636, "train_token_loss": 0.6679934723992095, "train_fea_loss": 1.0713579858914555, "epoch": 4, "n_parameters": 144845568}
+{"train_lr": 1.0000000000000597e-06, "train_loss": 5.909162417715259, "train_task_loss": 3.1000390200824333, "train_bpp_loss": 28.091233481177323, "train_patch_loss": 1.2870068884135293, "train_token_loss": 0.6689227937365607, "train_fea_loss": 1.1441093232500539, "epoch": 0, "n_parameters": 144845568}
+{"train_lr": 1.0000000000000597e-06, "train_loss": 5.4208089515959905, "train_task_loss": 2.6417781315404447, "train_bpp_loss": 27.790307712640693, "train_patch_loss": 1.0698811880355295, "train_token_loss": 0.6684176272018064, "train_fea_loss": 0.9034793063801684, "epoch": 1, "n_parameters": 144845568}
+{"train_lr": 2.330000000000072e-05, "train_loss": 4.90253949746025, "train_task_loss": 2.468064440561713, "train_bpp_loss": 24.344750106834965, "train_patch_loss": 0.9905159351565569, "train_token_loss": 0.6681102167086885, "train_fea_loss": 0.809438281672464, "epoch": 2, "n_parameters": 144845568}
+{"train_lr": 4.5600000000004605e-05, "train_loss": 4.039170976486995, "train_task_loss": 2.4127648418648637, "train_bpp_loss": 16.264061019539263, "train_patch_loss": 0.9645157243489183, "train_token_loss": 0.668388756130525, "train_fea_loss": 0.7798603551274688, "epoch": 3, "n_parameters": 144845568}
+{"train_lr": 6.790000000000497e-05, "train_loss": 3.154224511977437, "train_task_loss": 2.4628324630252605, "train_bpp_loss": 6.913920321606761, "train_patch_loss": 0.9734525120118075, "train_token_loss": 0.668419569109877, "train_fea_loss": 0.8209603751900134, "epoch": 4, "n_parameters": 144845568}
+{"train_lr": 9.019999999999779e-05, "train_loss": 2.855963341034145, "train_task_loss": 2.595376729947343, "train_bpp_loss": 2.605866110010399, "train_patch_loss": 1.0060626347426422, "train_token_loss": 0.7175173823357665, "train_fea_loss": 0.8717967045071314, "epoch": 5, "n_parameters": 144845568}
+{"train_lr": 0.00011234201335381617, "train_loss": 2.8053760061786472, "train_task_loss": 2.636510731433507, "train_bpp_loss": 1.6886527469109216, "train_patch_loss": 1.019217302998175, "train_token_loss": 0.7347769178539443, "train_fea_loss": 0.8825165003926742, "epoch": 6, "n_parameters": 144845568}
+{"train_lr": 0.00011227255068590994, "train_loss": 2.761493076880773, "train_task_loss": 2.6242000526464957, "train_bpp_loss": 1.3729302407391994, "train_patch_loss": 1.0161546610956855, "train_token_loss": 0.735874281325143, "train_fea_loss": 0.8721711006757024, "epoch": 7, "n_parameters": 144845568}
+{"train_lr": 0.0001121904989670886, "train_loss": 2.7252962913819783, "train_task_loss": 2.6068851682994008, "train_bpp_loss": 1.1841112330144972, "train_patch_loss": 1.01152529233812, "train_token_loss": 0.7345748117854662, "train_fea_loss": 0.8607850549831915, "epoch": 8, "n_parameters": 144845568}
+{"train_lr": 0.00011209587844235662, "train_loss": 2.6992331013548, "train_task_loss": 2.594223390594661, "train_bpp_loss": 1.0500971039855402, "train_patch_loss": 1.00772097030525, "train_token_loss": 0.7340729326429627, "train_fea_loss": 0.8524294771026257, "epoch": 9, "n_parameters": 144845568}
+{"train_lr": 0.0001119887124579783, "train_loss": 2.678213749384637, "train_task_loss": 2.5821204649494423, "train_bpp_loss": 0.9609328405931592, "train_patch_loss": 1.0029366761982013, "train_token_loss": 0.7342181133186753, "train_fea_loss": 0.8449656662421761, "epoch": 10, "n_parameters": 144845568}
+{"train_lr": 0.00011186902745551124, "train_loss": 2.662179193613555, "train_task_loss": 2.571508611331312, "train_bpp_loss": 0.9067058223556522, "train_patch_loss": 0.9990729609732147, "train_token_loss": 0.7330601164686987, "train_fea_loss": 0.8393755242634985, "epoch": 11, "n_parameters": 144845568}
+{"train_lr": 0.00011173685296543875, "train_loss": 2.6492657475530814, "train_task_loss": 2.562294816915437, "train_bpp_loss": 0.8697093047259702, "train_patch_loss": 0.995477003563526, "train_token_loss": 0.7320167214611916, "train_fea_loss": 0.834801082157617, "epoch": 12, "n_parameters": 144845568}
+{"train_lr": 0.00011159222159984347, "train_loss": 2.6404250215033263, "train_task_loss": 2.5562614543492987, "train_bpp_loss": 0.8416356684604316, "train_patch_loss": 0.9931515277579999, "train_token_loss": 0.7319663158244938, "train_fea_loss": 0.8311436018064546, "epoch": 13, "n_parameters": 144845568}
+{"train_lr": 0.00011143516904437194, "train_loss": 2.631879109013781, "train_task_loss": 2.5498513719625326, "train_bpp_loss": 0.8202773725337011, "train_patch_loss": 0.99033929754665, "train_token_loss": 0.7316777065275706, "train_fea_loss": 0.8278343591714006, "epoch": 14, "n_parameters": 144845568}
+{"train_lr": 0.00011126573404935182, "train_loss": 2.6259225274054265, "train_task_loss": 2.5454919715579467, "train_bpp_loss": 0.8043055602369441, "train_patch_loss": 0.9886904231225737, "train_token_loss": 0.731208170509542, "train_fea_loss": 0.8255933674758108, "epoch": 15, "n_parameters": 144845568}
+{"train_lr": 0.0001110839584203689, "train_loss": 2.6198472417533685, "train_task_loss": 2.5406655372529148, "train_bpp_loss": 0.791817048280913, "train_patch_loss": 0.9866770549277536, "train_token_loss": 0.7308970645976831, "train_fea_loss": 0.8230914076035328, "epoch": 16, "n_parameters": 144845568}
+{"train_lr": 0.0001108898870078482, "train_loss": 2.6147061569584933, "train_task_loss": 2.5367300099552534, "train_bpp_loss": 0.7797614707582468, "train_patch_loss": 0.985000457371221, "train_token_loss": 0.7312123553148455, "train_fea_loss": 0.8205171879849953, "epoch": 17, "n_parameters": 144845568}
+{"train_lr": 0.00011068356769595686, "train_loss": 2.609768410336128, "train_task_loss": 2.5328920738910026, "train_bpp_loss": 0.7687633632072549, "train_patch_loss": 0.9837793158322609, "train_token_loss": 0.7307359146476864, "train_fea_loss": 0.8183768322217844, "epoch": 18, "n_parameters": 144845568}
+{"train_lr": 0.0001104650513909484, "train_loss": 2.6053047203313677, "train_task_loss": 2.52939695734486, "train_bpp_loss": 0.7590776344502179, "train_patch_loss": 0.9826806558796363, "train_token_loss": 0.7301465437459431, "train_fea_loss": 0.8165697485534585, "epoch": 19, "n_parameters": 144845568}
+{"train_lr": 0.00011023439200841275, "train_loss": 2.60126930419847, "train_task_loss": 2.526219468670998, "train_bpp_loss": 0.7504983498341888, "train_patch_loss": 0.981161863478325, "train_token_loss": 0.7301577162488901, "train_fea_loss": 0.8148998804368096, "epoch": 20, "n_parameters": 144845568}
+{"train_lr": 0.00010973687463990103, "train_loss": 2.6637201429449684, "train_task_loss": 2.5465192173382074, "train_bpp_loss": 0.5860046291750648, "train_patch_loss": 0.9867626110592084, "train_token_loss": 0.7298668318939545, "train_fea_loss": 0.8298897645157054, "epoch": 21, "n_parameters": 144845568}
+{"train_lr": 0.00010973687463990103, "train_loss": 2.6604584804970584, "train_task_loss": 2.5488337897550433, "train_bpp_loss": 0.5581234564080084, "train_patch_loss": 0.98777336684643, "train_token_loss": 0.730239200644749, "train_fea_loss": 0.8308212128188684, "epoch": 22, "n_parameters": 144845568}
+{"train_lr": 0.00010947013940891518, "train_loss": 2.6580159922512316, "train_task_loss": 2.5488927058553954, "train_bpp_loss": 0.5456164319233464, "train_patch_loss": 0.9876678835765885, "train_token_loss": 0.7302756367019004, "train_fea_loss": 0.8309491750956189, "epoch": 23, "n_parameters": 144845568}
+{"train_lr": 0.00010919150658002209, "train_loss": 2.6552707054864446, "train_task_loss": 2.548027373189263, "train_bpp_loss": 0.5362166606305208, "train_patch_loss": 0.987526344177838, "train_token_loss": 0.7304425147100747, "train_fea_loss": 0.830058504461045, "epoch": 24, "n_parameters": 144845568}
+{"train_lr": 0.00010890104490178482, "train_loss": 2.6532317997365116, "train_task_loss": 2.5475874447243676, "train_bpp_loss": 0.5282217790852848, "train_patch_loss": 0.9868854832560788, "train_token_loss": 0.7307839357099838, "train_fea_loss": 0.8299180160475363, "epoch": 25, "n_parameters": 144845568}
+{"train_lr": 0.00010828492456631287, "train_loss": 2.6756018440750577, "train_task_loss": 2.554671129796931, "train_bpp_loss": 0.48372285632077133, "train_patch_loss": 0.9893747716762429, "train_token_loss": 0.7302002353723244, "train_fea_loss": 0.83509611297506, "epoch": 26, "n_parameters": 144845568}
+{"train_lr": 0.00010828492456631287, "train_loss": 2.672080845176745, "train_task_loss": 2.5541286950649544, "train_bpp_loss": 0.4718086026637978, "train_patch_loss": 0.9892294842401426, "train_token_loss": 0.7304161941873835, "train_fea_loss": 0.8344830075631754, "epoch": 27, "n_parameters": 144845568}
+{"train_lr": 0.00010795941792757138, "train_loss": 2.6682109270039365, "train_task_loss": 2.552076613383113, "train_bpp_loss": 0.4645372550583911, "train_patch_loss": 0.9893483498738967, "train_token_loss": 0.7297189640332398, "train_fea_loss": 0.8330092908327671, "epoch": 28, "n_parameters": 144845568}
+{"train_lr": 0.00010762238643889585, "train_loss": 2.6648726280626778, "train_task_loss": 2.5502606650068462, "train_bpp_loss": 0.458447851656045, "train_patch_loss": 0.988596663794679, "train_token_loss": 0.7296670780559595, "train_fea_loss": 0.8319969125855848, "epoch": 29, "n_parameters": 144845568}
+{"train_lr": 0.00010727391325772412, "train_loss": 2.66122928966614, "train_task_loss": 2.5479887132521726, "train_bpp_loss": 0.4529623082800985, "train_patch_loss": 0.9877295848607517, "train_token_loss": 0.7298175230917099, "train_fea_loss": 0.8304415963706281, "epoch": 30, "n_parameters": 144845568}
+{"train_lr": 0.00010691408436465084, "train_loss": 2.6585902343425962, "train_task_loss": 2.546459772737978, "train_bpp_loss": 0.44852184575292275, "train_patch_loss": 0.9874212175062734, "train_token_loss": 0.7292956591402884, "train_fea_loss": 0.8297428859540705, "epoch": 31, "n_parameters": 144845568}
+{"train_lr": 0.00010654298854205488, "train_loss": 2.655904775418395, "train_task_loss": 2.5447790948875086, "train_bpp_loss": 0.44450271940216135, "train_patch_loss": 0.9865645243444567, "train_token_loss": 0.7294259656095605, "train_fea_loss": 0.8287885952718573, "epoch": 32, "n_parameters": 144845568}
+{"train_lr": 0.0001061607173522522, "train_loss": 2.653865991727554, "train_task_loss": 2.5435698100184796, "train_bpp_loss": 0.4411847302678981, "train_patch_loss": 0.98615532235535, "train_token_loss": 0.7293780767236003, "train_fea_loss": 0.8280364017178722, "epoch": 33, "n_parameters": 144845568}
+{"train_lr": 0.00010576736511496153, "train_loss": 2.649993774818359, "train_task_loss": 2.540434418590092, "train_bpp_loss": 0.4382374218332289, "train_patch_loss": 0.9851716845977506, "train_token_loss": 0.7283954293907082, "train_fea_loss": 0.8268672945759565, "epoch": 34, "n_parameters": 144845568}
+{"train_lr": 0.00010536302888396166, "train_loss": 2.64796845856116, "train_task_loss": 2.539067151297411, "train_bpp_loss": 0.4356052296792271, "train_patch_loss": 0.9849777258128571, "train_token_loss": 0.7279274276853036, "train_fea_loss": 0.8261619895200828, "epoch": 35, "n_parameters": 144845568}
+{"train_lr": 0.00010494780842314013, "train_loss": 2.646916040002728, "train_task_loss": 2.538628663867712, "train_bpp_loss": 0.43314950562102333, "train_patch_loss": 0.9843870765022594, "train_token_loss": 0.7287507657747831, "train_fea_loss": 0.8254908122640434, "epoch": 36, "n_parameters": 144845568}
+{"train_lr": 0.00010452180618197858, "train_loss": 2.6438271829979025, "train_task_loss": 2.5360889318678304, "train_bpp_loss": 0.43095300335781, "train_patch_loss": 0.9832675016509722, "train_token_loss": 0.7280113412965116, "train_fea_loss": 0.8248100794753511, "epoch": 37, "n_parameters": 144845568}
+{"train_lr": 0.00010408512727011787, "train_loss": 2.6436403147715457, "train_task_loss": 2.536438628754241, "train_bpp_loss": 0.42880674428697363, "train_patch_loss": 0.9832872114673579, "train_token_loss": 0.7284368306166119, "train_fea_loss": 0.8247145772900442, "epoch": 38, "n_parameters": 144845568}
+{"train_lr": 0.00010363787943157281, "train_loss": 2.640265642149414, "train_task_loss": 2.5335381660583636, "train_bpp_loss": 0.42690990286783204, "train_patch_loss": 0.982539902213398, "train_token_loss": 0.7274603552767913, "train_fea_loss": 0.8235378990137427, "epoch": 39, "n_parameters": 144845568}
+{"train_lr": 0.0001031801730180277, "train_loss": 2.638479518384742, "train_task_loss": 2.532214540898871, "train_bpp_loss": 0.4250599086704514, "train_patch_loss": 0.9822922450729268, "train_token_loss": 0.7268436730546738, "train_fea_loss": 0.8230786147199088, "epoch": 40, "n_parameters": 144845568}
+{"train_lr": 0.00010271212096170505, "train_loss": 2.6368867533473517, "train_task_loss": 2.5310158089464374, "train_bpp_loss": 0.4234837778348127, "train_patch_loss": 0.9816671004675204, "train_token_loss": 0.7268876022875881, "train_fea_loss": 0.8224610961522523, "epoch": 41, "n_parameters": 144845568}
+{"train_lr": 0.00010223383874746677, "train_loss": 2.6354142183826554, "train_task_loss": 2.529919751140354, "train_bpp_loss": 0.42197786968612344, "train_patch_loss": 0.9807341171861839, "train_token_loss": 0.7272372794665879, "train_fea_loss": 0.8219483458689398, "epoch": 42, "n_parameters": 144845568}
+{"train_lr": 0.00010174544438424974, "train_loss": 2.634054534628594, "train_task_loss": 2.528918670253645, "train_bpp_loss": 0.4205434575589816, "train_patch_loss": 0.9803530661996713, "train_token_loss": 0.7265599571498964, "train_fea_loss": 0.8220056383613702, "epoch": 43, "n_parameters": 144845568}
+{"train_lr": 0.00010124705837609591, "train_loss": 2.6349923593892184, "train_task_loss": 2.5301807783186723, "train_bpp_loss": 0.4192463266741422, "train_patch_loss": 0.9804664164607366, "train_token_loss": 0.726979716709978, "train_fea_loss": 0.8227346358883331, "epoch": 44, "n_parameters": 144845568}
+{"train_lr": 0.00010073880369226542, "train_loss": 2.63449551127583, "train_task_loss": 2.529952161675163, "train_bpp_loss": 0.4181733996109318, "train_patch_loss": 0.9800947665212323, "train_token_loss": 0.7259607360401242, "train_fea_loss": 0.8238966495861753, "epoch": 45, "n_parameters": 144845568}
+{"train_lr": 0.00010022080573700511, "train_loss": 2.6345730214078222, "train_task_loss": 2.530297286192076, "train_bpp_loss": 0.4171029411935263, "train_patch_loss": 0.9800362460766372, "train_token_loss": 0.7260628667832029, "train_fea_loss": 0.8241981637554703, "epoch": 46, "n_parameters": 144845568}
+{"train_lr": 9.969319231856176e-05, "train_loss": 2.6355131778874985, "train_task_loss": 2.5314376252464394, "train_bpp_loss": 0.4163022163889105, "train_patch_loss": 0.9805744725929426, "train_token_loss": 0.7263783447029029, "train_fea_loss": 0.8244847989500427, "epoch": 47, "n_parameters": 144845568}
+{"train_lr": 9.915609361765753e-05, "train_loss": 2.634658775285637, "train_task_loss": 2.5308349787099043, "train_bpp_loss": 0.4152951848266615, "train_patch_loss": 0.9807003051248029, "train_token_loss": 0.725708744122381, "train_fea_loss": 0.8244259209490473, "epoch": 48, "n_parameters": 144845568}
+{"train_lr": 9.860964215535301e-05, "train_loss": 2.6395780852765776, "train_task_loss": 2.536209069951761, "train_bpp_loss": 0.4134760623657071, "train_patch_loss": 0.9824451866852258, "train_token_loss": 0.7274219641725508, "train_fea_loss": 0.8263419094569177, "epoch": 49, "n_parameters": 144845568}
+{"train_lr": 9.805397276035986e-05, "train_loss": 2.639889271710988, "train_task_loss": 2.536885856259927, "train_bpp_loss": 0.4120136631007264, "train_patch_loss": 0.9825187112367768, "train_token_loss": 0.7275711355386949, "train_fea_loss": 0.8267959996776508, "epoch": 50, "n_parameters": 144845568}
+{"train_lr": 9.748922253581646e-05, "train_loss": 2.6434670704320893, "train_task_loss": 2.5408570495137064, "train_bpp_loss": 0.4104400824474023, "train_patch_loss": 0.9840436073498046, "train_token_loss": 0.7284569417749592, "train_fea_loss": 0.828356491024349, "epoch": 51, "n_parameters": 144845568}
+{"train_lr": 9.691553082535863e-05, "train_loss": 2.6544609911453954, "train_task_loss": 2.5523642709644005, "train_bpp_loss": 0.40838687800259577, "train_patch_loss": 0.9880639841150252, "train_token_loss": 0.7311584631057023, "train_fea_loss": 0.8331418140172708, "epoch": 52, "n_parameters": 144845568}
+{"train_lr": 9.633303917884302e-05, "train_loss": 2.6421383294395264, "train_task_loss": 2.5401510994175642, "train_bpp_loss": 0.40794892040498965, "train_patch_loss": 0.9832984244683199, "train_token_loss": 0.7280591241032671, "train_fea_loss": 0.8287935424283397, "epoch": 53, "n_parameters": 144845568}
+{"train_lr": 9.574189131737902e-05, "train_loss": 2.6519711823504175, "train_task_loss": 2.5503779566384477, "train_bpp_loss": 0.40637290150338085, "train_patch_loss": 0.987097349643207, "train_token_loss": 0.7309697274904814, "train_fea_loss": 0.832310870747647, "epoch": 54, "n_parameters": 144845568}
+{"train_lr": 9.514223309782753e-05, "train_loss": 2.649527451563939, "train_task_loss": 2.5482009406730852, "train_bpp_loss": 0.4053060460971932, "train_patch_loss": 0.9869010643251865, "train_token_loss": 0.7299437226093323, "train_fea_loss": 0.8313561448876521, "epoch": 55, "n_parameters": 144845568}
+{"train_lr": 9.453421247691757e-05, "train_loss": 2.643370286264508, "train_task_loss": 2.5421072299049485, "train_bpp_loss": 0.40505222469563784, "train_patch_loss": 0.9843411268666387, "train_token_loss": 0.7289866340465004, "train_fea_loss": 0.8287794603441331, "epoch": 56, "n_parameters": 144845568}
+{"train_lr": 9.391797947461475e-05, "train_loss": 2.6461283357577217, "train_task_loss": 2.5451179012316736, "train_bpp_loss": 0.40404173961842804, "train_patch_loss": 0.9857150221699111, "train_token_loss": 0.7291283701508523, "train_fea_loss": 0.8302744993274447, "epoch": 57, "n_parameters": 144845568}
+{"train_lr": 9.329368613720009e-05, "train_loss": 2.652468080438084, "train_task_loss": 2.551866788899513, "train_bpp_loss": 0.40240516668471465, "train_patch_loss": 0.9880323676700643, "train_token_loss": 0.7304665103925468, "train_fea_loss": 0.8333679016130112, "epoch": 58, "n_parameters": 144845568}
+{"train_lr": 9.266148649972007e-05, "train_loss": 2.650359462124767, "train_task_loss": 2.5499309687752376, "train_bpp_loss": 0.401713975702973, "train_patch_loss": 0.9875440163452991, "train_token_loss": 0.7306666403307581, "train_fea_loss": 0.8317203026117502, "epoch": 59, "n_parameters": 144845568}
+{"train_lr": 9.202153654795684e-05, "train_loss": 2.6472997942613326, "train_task_loss": 2.547006531638636, "train_bpp_loss": 0.4011730529006115, "train_patch_loss": 0.986042402963862, "train_token_loss": 0.7303420982616161, "train_fea_loss": 0.8306220213389046, "epoch": 60, "n_parameters": 144845568}
+{"train_lr": 9.137399417998249e-05, "train_loss": 2.635524819071273, "train_task_loss": 2.5351910238875615, "train_bpp_loss": 0.40133518058520784, "train_patch_loss": 0.9820189805460169, "train_token_loss": 0.7271724717727275, "train_fea_loss": 0.8259995626531345, "epoch": 61, "n_parameters": 144845568}
+{"train_lr": 9.071901916722404e-05, "train_loss": 2.648858120589019, "train_task_loss": 2.5489517275056395, "train_bpp_loss": 0.39962557505043467, "train_patch_loss": 0.986155892487803, "train_token_loss": 0.73087697630029, "train_fea_loss": 0.8319188496388262, "epoch": 62, "n_parameters": 144845568}
+{"train_lr": 9.071901916722404e-05, "train_loss": 2.7465927910497436, "train_task_loss": 2.5472986304449092, "train_bpp_loss": 0.3870172494111194, "train_patch_loss": 0.986047504940997, "train_token_loss": 0.7293615024804855, "train_fea_loss": 0.9150785957809261, "epoch": 61, "n_parameters": 144845568}
+{"train_lr": 9.071901916722404e-05, "train_loss": 2.732339255324156, "train_task_loss": 2.5342044824247454, "train_bpp_loss": 0.3851228334161518, "train_patch_loss": 0.9823006979486854, "train_token_loss": 0.725924804654839, "train_fea_loss": 0.9085768888921117, "epoch": 62, "n_parameters": 144845568}
+{"train_lr": 9.005677311491453e-05, "train_loss": 2.733937658449943, "train_task_loss": 2.5361717377456543, "train_bpp_loss": 0.38358666918550927, "train_patch_loss": 0.9829422572395391, "train_token_loss": 0.7263305272473741, "train_fea_loss": 0.9095888588916984, "epoch": 63, "n_parameters": 144845568}
+{"train_lr": 8.938741942239847e-05, "train_loss": 2.735222797715764, "train_task_loss": 2.5376820612332516, "train_bpp_loss": 0.382704222326924, "train_patch_loss": 0.9836816218920605, "train_token_loss": 0.7267059973671103, "train_fea_loss": 0.9100238969839877, "epoch": 64, "n_parameters": 144845568}
+{"train_lr": 8.871112324267081e-05, "train_loss": 2.7348346617742836, "train_task_loss": 2.537643769224413, "train_bpp_loss": 0.38169531316588, "train_patch_loss": 0.9843459360355096, "train_token_loss": 0.726475094608087, "train_fea_loss": 0.9095050229219724, "epoch": 65, "n_parameters": 144845568}
+{"train_lr": 8.733837255720078e-05, "train_loss": 2.7936204858570934, "train_task_loss": 2.558094645236179, "train_bpp_loss": 0.3368677387395941, "train_patch_loss": 0.9911075603882435, "train_token_loss": 0.7276336732755998, "train_fea_loss": 0.9232887631746755, "epoch": 66, "n_parameters": 144845568}
+{"train_lr": 8.733837255720078e-05, "train_loss": 2.7915395827209064, "train_task_loss": 2.558850450686914, "train_bpp_loss": 0.3302956995413565, "train_patch_loss": 0.9912216057375574, "train_token_loss": 0.7270683532455664, "train_fea_loss": 0.9246165518954896, "epoch": 67, "n_parameters": 144845568}
+{"train_lr": 8.66422567571558e-05, "train_loss": 2.8010428122443547, "train_task_loss": 2.569311464486791, "train_bpp_loss": 0.3271386409915394, "train_patch_loss": 0.9944994980447012, "train_token_loss": 0.729622550842484, "train_fea_loss": 0.929708367548615, "epoch": 68, "n_parameters": 144845568}
+{"train_lr": 8.59398757977085e-05, "train_loss": 2.8097731665634424, "train_task_loss": 2.5788058781241485, "train_bpp_loss": 0.3244854515335775, "train_patch_loss": 0.9977880177184117, "train_token_loss": 0.7315296879824367, "train_fea_loss": 0.9344370006996807, "epoch": 69, "n_parameters": 144845568}
+{"train_lr": 8.523140298084917e-05, "train_loss": 2.7927491063747905, "train_task_loss": 2.562695958050821, "train_bpp_loss": 0.3240329879190889, "train_patch_loss": 0.9920173354035945, "train_token_loss": 0.7282957452923345, "train_fea_loss": 0.9266211757105853, "epoch": 70, "n_parameters": 144845568}
+{"train_lr": 8.451701311164659e-05, "train_loss": 2.816440251421371, "train_task_loss": 2.5866508388297733, "train_bpp_loss": 0.32121099593403707, "train_patch_loss": 1.0007479091105667, "train_token_loss": 0.7334584623767282, "train_fea_loss": 0.9376889240237878, "epoch": 71, "n_parameters": 144845568}
+{"train_lr": 8.379688245511898e-05, "train_loss": 2.8120930993692768, "train_task_loss": 2.5828994919570514, "train_bpp_loss": 0.32026763344004894, "train_patch_loss": 0.9992939126554975, "train_token_loss": 0.732874036779989, "train_fea_loss": 0.9358047071388729, "epoch": 72, "n_parameters": 144845568}
+{"train_lr": 8.307118869271464e-05, "train_loss": 2.8027984519763818, "train_task_loss": 2.5742826347978567, "train_bpp_loss": 0.3194215351233272, "train_patch_loss": 0.9962847879741618, "train_token_loss": 0.7302367620156716, "train_fea_loss": 0.932537203840018, "epoch": 73, "n_parameters": 144845568}
+{"train_lr": 8.234011087850579e-05, "train_loss": 2.832480542442138, "train_task_loss": 2.6040796041131307, "train_bpp_loss": 0.3165612864526931, "train_patch_loss": 1.0063268801385015, "train_token_loss": 0.7382693186433589, "train_fea_loss": 0.945431755996383, "epoch": 74, "n_parameters": 144845568}
+{"train_lr": 8.160382939503717e-05, "train_loss": 2.799674530850826, "train_task_loss": 2.5721579685592824, "train_bpp_loss": 0.3175276949637674, "train_patch_loss": 0.9966744806246756, "train_token_loss": 0.7291926598819576, "train_fea_loss": 0.9309199213762757, "epoch": 75, "n_parameters": 144845568}
+{"train_lr": 8.011638332509435e-05, "train_loss": 2.8255829363078666, "train_task_loss": 2.578499776871799, "train_bpp_loss": 0.2999681402433315, "train_patch_loss": 0.9987926142361703, "train_token_loss": 0.7287038154274725, "train_fea_loss": 0.9361036925038381, "epoch": 76, "n_parameters": 144845568}
+{"train_lr": 8.011638332509435e-05, "train_loss": 2.840959037218591, "train_task_loss": 2.5953164593778926, "train_bpp_loss": 0.2960671754754919, "train_patch_loss": 1.0047289559068224, "train_token_loss": 0.7329247621774709, "train_fea_loss": 0.9434290263216666, "epoch": 77, "n_parameters": 144845568}
+{"train_lr": 7.93655857436786e-05, "train_loss": 2.8381004938922767, "train_task_loss": 2.5934427621946345, "train_bpp_loss": 0.29446334153194903, "train_patch_loss": 1.0050567829807242, "train_token_loss": 0.7319110037146999, "train_fea_loss": 0.9421224841419206, "epoch": 78, "n_parameters": 144845568}
+{"train_lr": 7.86103184125689e-05, "train_loss": 2.847826922665254, "train_task_loss": 2.603588223600273, "train_bpp_loss": 0.29271904649034286, "train_patch_loss": 1.0064198724385038, "train_token_loss": 0.7354645022689152, "train_fea_loss": 0.9478742439049682, "epoch": 79, "n_parameters": 144845568}
+{"train_lr": 7.785076768264985e-05, "train_loss": 2.8496559623369784, "train_task_loss": 2.6059850953447876, "train_bpp_loss": 0.29151453778183434, "train_patch_loss": 1.0084720284899147, "train_token_loss": 0.7349831929711772, "train_fea_loss": 0.9487828727950045, "epoch": 80, "n_parameters": 144845568}
+{"train_lr": 7.708712096171631e-05, "train_loss": 2.8463599768867023, "train_task_loss": 2.6032515673292913, "train_bpp_loss": 0.2907897324716712, "train_patch_loss": 1.0084952734437564, "train_token_loss": 0.7339370518864737, "train_fea_loss": 0.9469011765759554, "epoch": 81, "n_parameters": 144845568}
+{"train_lr": 7.631956666815207e-05, "train_loss": 2.8331575906015845, "train_task_loss": 2.5902766323418356, "train_bpp_loss": 0.29128668203445185, "train_patch_loss": 1.0040304742896442, "train_token_loss": 0.7303849485070764, "train_fea_loss": 0.9414473412786242, "epoch": 82, "n_parameters": 144845568}
+{"train_lr": 7.554829418450765e-05, "train_loss": 6.022797273306681, "train_task_loss": 5.747255473551776, "train_bpp_loss": 0.15595893754092843, "train_patch_loss": 1.9295167815905037, "train_token_loss": 1.9044995418676047, "train_fea_loss": 2.1045631034237196, "epoch": 83, "n_parameters": 144845568}
+{"train_lr": 7.477349381072652e-05, "train_loss": 6.33198429772751, "train_task_loss": 6.103147584757359, "train_bpp_loss": 0.04776455266578447, "train_patch_loss": 2.038426331773722, "train_token_loss": 2.034283296100134, "train_fea_loss": 2.2334817904058837, "epoch": 84, "n_parameters": 144845568}
+{"train_lr": 7.399535671720344e-05, "train_loss": 6.269092192288211, "train_task_loss": 6.053333911088874, "train_bpp_loss": 0.029455166824030238, "train_patch_loss": 2.0307816479676704, "train_token_loss": 2.024027929974081, "train_fea_loss": 2.1983768022049675, "epoch": 85, "n_parameters": 144845568}
+{"train_lr": 7.321407489761549e-05, "train_loss": 6.215953613303119, "train_task_loss": 6.003636568552442, "train_bpp_loss": 0.026147063460792316, "train_patch_loss": 2.018133198333194, "train_token_loss": 2.0035276831500677, "train_fea_loss": 2.1801732855288387, "epoch": 86, "n_parameters": 144845568}
+{"train_lr": 7.242984112156774e-05, "train_loss": 5.868365104446451, "train_task_loss": 5.527729460429088, "train_bpp_loss": 0.2684132926173316, "train_patch_loss": 1.9746971274469254, "train_token_loss": 1.5961082754723674, "train_fea_loss": 2.1526165047361197, "epoch": 87, "n_parameters": 144845568}
+{"train_lr": 7.16428488870196e-05, "train_loss": 5.4557995484958735, "train_task_loss": 5.110851676522685, "train_bpp_loss": 0.2943931522651305, "train_patch_loss": 1.8687771737611265, "train_token_loss": 1.3823193566252787, "train_fea_loss": 2.0457307090067105, "epoch": 88, "n_parameters": 144845568}
+{"train_lr": 7.085329237251759e-05, "train_loss": 6.162637532811513, "train_task_loss": 5.88701508848144, "train_bpp_loss": 0.13411895593832643, "train_patch_loss": 2.0334672068672286, "train_token_loss": 1.821566376444265, "train_fea_loss": 2.2351796914660316, "epoch": 89, "n_parameters": 144845568}
+{"train_lr": 7.006136638931818e-05, "train_loss": 5.024170764112215, "train_task_loss": 4.809058998057025, "train_bpp_loss": 0.05317469353166999, "train_patch_loss": 1.8792113221967393, "train_token_loss": 1.0658738444262235, "train_fea_loss": 2.0503712694118206, "epoch": 90, "n_parameters": 144845568}
+{"train_lr": 6.926726633331106e-05, "train_loss": 4.723379994860942, "train_task_loss": 4.515086256485763, "train_bpp_loss": 0.057305316362173765, "train_patch_loss": 1.798007144451999, "train_token_loss": 0.9435908788395264, "train_fea_loss": 1.9508371142764314, "epoch": 91, "n_parameters": 144845568}
+{"train_lr": 6.847118813679865e-05, "train_loss": 5.031990508435013, "train_task_loss": 4.765979826932759, "train_bpp_loss": 0.15900398099812293, "train_patch_loss": 1.824176159184828, "train_token_loss": 1.1403188319715987, "train_fea_loss": 1.981633366195025, "epoch": 92, "n_parameters": 144845568}
+{"train_lr": 7.631956666815207e-05, "train_loss": 2.814900835471259, "train_task_loss": 2.5911692696259463, "train_bpp_loss": 0.30722877928625336, "train_patch_loss": 1.003347651021247, "train_token_loss": 0.7330356562840674, "train_fea_loss": 0.854785952643364, "epoch": 81, "n_parameters": 144845568}
+{"train_lr": 7.631956666815207e-05, "train_loss": 2.8142806054126446, "train_task_loss": 2.5897073253131597, "train_bpp_loss": 0.30940049388993096, "train_patch_loss": 1.0041070785507453, "train_token_loss": 0.7321698584502263, "train_fea_loss": 0.8534303789900289, "epoch": 82, "n_parameters": 144845568}
+{"train_lr": 7.554829418450765e-05, "train_loss": 2.816245504384704, "train_task_loss": 2.591411675659301, "train_bpp_loss": 0.30987349456864044, "train_patch_loss": 1.003780739015634, "train_token_loss": 0.7337235686563545, "train_fea_loss": 0.8539073579683364, "epoch": 83, "n_parameters": 144845568}
+{"train_lr": 7.477349381072652e-05, "train_loss": 2.800685787649392, "train_task_loss": 2.5759294760688176, "train_bpp_loss": 0.3108195169511691, "train_patch_loss": 0.9979896100653376, "train_token_loss": 0.729064737578376, "train_fea_loss": 0.8488751188081374, "epoch": 84, "n_parameters": 144845568}
+{"train_lr": 7.399535671720344e-05, "train_loss": 2.78943671736357, "train_task_loss": 2.5648709468239788, "train_bpp_loss": 0.31156991790573624, "train_patch_loss": 0.9950137249611729, "train_token_loss": 0.7262643160526272, "train_fea_loss": 0.8435928968185608, "epoch": 85, "n_parameters": 144845568}
+{"train_lr": 7.321407489761549e-05, "train_loss": 2.844264631443244, "train_task_loss": 2.614422207854563, "train_bpp_loss": 0.319169871819516, "train_patch_loss": 1.0113446999082176, "train_token_loss": 0.7409178713242784, "train_fea_loss": 0.8621596260023096, "epoch": 86, "n_parameters": 144845568}
+{"train_lr": 7.242984112156774e-05, "train_loss": 2.851643653814312, "train_task_loss": 2.6122646205198707, "train_bpp_loss": 0.33907441596994153, "train_patch_loss": 1.01671234990651, "train_token_loss": 0.7275969981968545, "train_fea_loss": 0.8679552621183564, "epoch": 87, "n_parameters": 144845568}
+{"train_lr": 7.16428488870196e-05, "train_loss": 2.778331270660285, "train_task_loss": 2.552063302813674, "train_bpp_loss": 0.31613819805538945, "train_patch_loss": 0.9926303882571719, "train_token_loss": 0.7193752784768848, "train_fea_loss": 0.8400576262571793, "epoch": 88, "n_parameters": 144845568}
+{"train_lr": 7.085329237251759e-05, "train_loss": 2.7897735733261926, "train_task_loss": 2.565236364247845, "train_bpp_loss": 0.311233962220486, "train_patch_loss": 0.9957534326776434, "train_token_loss": 0.7246638578329262, "train_fea_loss": 0.8448190648729602, "epoch": 89, "n_parameters": 144845568}
+{"train_lr": 7.006136638931818e-05, "train_loss": 2.7982831528093173, "train_task_loss": 2.574218000189292, "train_bpp_loss": 0.3094736575059664, "train_patch_loss": 1.000583418184887, "train_token_loss": 0.7256146864831555, "train_fea_loss": 0.8480198857010447, "epoch": 90, "n_parameters": 144845568}
+{"train_lr": 6.926726633331106e-05, "train_loss": 3.415653905788843, "train_task_loss": 2.975828631783275, "train_bpp_loss": 0.7560260864588872, "train_patch_loss": 1.1288673198967278, "train_token_loss": 0.8508261374780814, "train_fea_loss": 0.9961351647445981, "epoch": 91, "n_parameters": 144845568}
+{"train_lr": 6.847118813679865e-05, "train_loss": 6.568326573541982, "train_task_loss": 6.1273969157779815, "train_bpp_loss": 0.5257196640345934, "train_patch_loss": 2.043614938980241, "train_token_loss": 2.040224435022302, "train_fea_loss": 2.04355752422548, "epoch": 92, "n_parameters": 144845568}
+{"train_lr": 6.767332822016792e-05, "train_loss": 6.867288129757062, "train_task_loss": 6.089932662566646, "train_bpp_loss": 1.2734335873903297, "train_patch_loss": 2.0432088341144063, "train_token_loss": 2.003620763126162, "train_fea_loss": 2.043103043592984, "epoch": 93, "n_parameters": 144845568}
+{"train_lr": 6.687388344341571e-05, "train_loss": 6.613793869366606, "train_task_loss": 4.8460818514299335, "train_bpp_loss": 3.475327093109524, "train_patch_loss": 2.043151048612859, "train_token_loss": 0.7647827833252167, "train_fea_loss": 2.038148040657671, "epoch": 94, "n_parameters": 144845568}
+{"train_lr": 6.607305105757049e-05, "train_loss": 6.583523838229174, "train_task_loss": 4.818476951540374, "train_bpp_loss": 3.4695032598624986, "train_patch_loss": 2.0428966481319004, "train_token_loss": 0.737876279511862, "train_fea_loss": 2.037704043197546, "epoch": 95, "n_parameters": 144845568}
+{"train_lr": 6.5271028656055e-05, "train_loss": 6.5769114312794, "train_task_loss": 4.8131763801979215, "train_bpp_loss": 3.46658337029586, "train_patch_loss": 2.0428622997046517, "train_token_loss": 0.7325888925904243, "train_fea_loss": 2.0377252054997057, "epoch": 96, "n_parameters": 144845568}
+{"train_lr": 6.446801412587525e-05, "train_loss": 6.572428462227329, "train_task_loss": 4.809816579738681, "train_bpp_loss": 3.4640789968754464, "train_patch_loss": 2.0429170189662087, "train_token_loss": 0.7291364202284806, "train_fea_loss": 2.0377631590699408, "epoch": 97, "n_parameters": 144845568}
+{"train_lr": 6.36642055988671e-05, "train_loss": 6.5666217358349614, "train_task_loss": 4.80531757448217, "train_bpp_loss": 3.461199431673443, "train_patch_loss": 2.0427514293746984, "train_token_loss": 0.7249221600507267, "train_fea_loss": 2.0376440027742078, "epoch": 98, "n_parameters": 144845568}
+{"train_lr": 6.285980140274965e-05, "train_loss": 6.492165406455668, "train_task_loss": 4.802894322652754, "train_bpp_loss": 3.3011300552377314, "train_patch_loss": 2.042857281175806, "train_token_loss": 0.72241166340549, "train_fea_loss": 2.0376253973210243, "epoch": 99, "n_parameters": 144845568}
+{"train_lr": 6.205500001222403e-05, "train_loss": 6.469163187473512, "train_task_loss": 4.799823627161751, "train_bpp_loss": 3.2568444883723338, "train_patch_loss": 2.042693032090255, "train_token_loss": 0.7195353948642953, "train_fea_loss": 2.0375952215226505, "epoch": 100, "n_parameters": 144845568}
+{"train_lr": 6.847118813679865e-05, "train_loss": 2.708930113305934, "train_task_loss": 2.5679492560668673, "train_bpp_loss": 0.3132908021231635, "train_patch_loss": 0.9976560541538062, "train_token_loss": 0.7208710700648723, "train_fea_loss": 0.8494221233278155, "epoch": 91, "n_parameters": 144845568}
+{"train_lr": 6.847118813679865e-05, "train_loss": 2.699149013628705, "train_task_loss": 2.5613451289723246, "train_bpp_loss": 0.3062308639035576, "train_patch_loss": 0.9944496202400978, "train_token_loss": 0.7207848698646724, "train_fea_loss": 0.846110629630264, "epoch": 92, "n_parameters": 144845568}
+{"train_lr": 6.767332822016792e-05, "train_loss": 2.695356560687867, "train_task_loss": 2.5580616601490433, "train_bpp_loss": 0.3050997874931805, "train_patch_loss": 0.9943495049304385, "train_token_loss": 0.7195412099754496, "train_fea_loss": 0.8441709351862113, "epoch": 93, "n_parameters": 144845568}
+{"train_lr": 6.687388344341571e-05, "train_loss": 2.709555460635921, "train_task_loss": 2.572864127852362, "train_bpp_loss": 0.3037585253364963, "train_patch_loss": 0.9997199533967306, "train_token_loss": 0.7228693755843835, "train_fea_loss": 0.850274788731669, "epoch": 94, "n_parameters": 144845568}
+{"train_lr": 6.607305105757049e-05, "train_loss": 2.7318896783841886, "train_task_loss": 2.59231706789965, "train_bpp_loss": 0.31016136522402615, "train_patch_loss": 1.0061367933327954, "train_token_loss": 0.7272278621984865, "train_fea_loss": 0.858952402501262, "epoch": 95, "n_parameters": 144845568}
+{"train_lr": 6.5271028656055e-05, "train_loss": 2.7111547499162545, "train_task_loss": 2.569625718514625, "train_bpp_loss": 0.31450896638912335, "train_patch_loss": 1.0005275200364543, "train_token_loss": 0.7153408098586207, "train_fea_loss": 0.8537573801527778, "epoch": 96, "n_parameters": 144845568}
+{"train_lr": 6.446801412587525e-05, "train_loss": 2.707958362782173, "train_task_loss": 2.570670375998, "train_bpp_loss": 0.30508442372374006, "train_patch_loss": 0.9985086321964752, "train_token_loss": 0.7216268923039809, "train_fea_loss": 0.8505348416952063, "epoch": 97, "n_parameters": 144845568}
+{"train_lr": 6.36642055988671e-05, "train_loss": 2.697063536604317, "train_task_loss": 2.560573679866265, "train_bpp_loss": 0.3033108016983014, "train_patch_loss": 0.9955314651368369, "train_token_loss": 0.7185008986127213, "train_fea_loss": 0.846541307080421, "epoch": 98, "n_parameters": 144845568}
+{"train_lr": 6.285980140274965e-05, "train_loss": 2.679979503681834, "train_task_loss": 2.5433403860768684, "train_bpp_loss": 0.30364249174190544, "train_patch_loss": 0.989279427558934, "train_token_loss": 0.7147660507581575, "train_fea_loss": 0.8392948986810568, "epoch": 99, "n_parameters": 144845568}
+{"train_lr": 6.205500001222403e-05, "train_loss": 2.818446820222145, "train_task_loss": 2.6755781511751584, "train_bpp_loss": 0.3174859403087635, "train_patch_loss": 1.0392651655052385, "train_token_loss": 0.7372288242098596, "train_fea_loss": 0.8990841494868699, "epoch": 100, "n_parameters": 144845568}
+{"train_lr": 6.12500000000064e-05, "train_loss": 2.7441475540310214, "train_task_loss": 2.597907278704629, "train_bpp_loss": 0.32497839667171025, "train_patch_loss": 1.0195179717119436, "train_token_loss": 0.7045413070298034, "train_fea_loss": 0.8738479904709853, "epoch": 101, "n_parameters": 144845568}
+{"train_lr": 6.044499998777186e-05, "train_loss": 2.7079060795299297, "train_task_loss": 2.565669994283137, "train_bpp_loss": 0.31608019693479905, "train_patch_loss": 1.0030313897252583, "train_token_loss": 0.7082743480091365, "train_fea_loss": 0.8543642482986089, "epoch": 102, "n_parameters": 144845568}
+{"train_lr": 5.964019859724661e-05, "train_loss": 2.696671325293519, "train_task_loss": 2.557040445006294, "train_bpp_loss": 0.3102908528876676, "train_patch_loss": 0.9976166626081157, "train_token_loss": 0.7102920439817922, "train_fea_loss": 0.8491317290227584, "epoch": 103, "n_parameters": 144845568}
+{"train_lr": 5.8835794401133974e-05, "train_loss": 2.691563926851578, "train_task_loss": 2.553079867954377, "train_bpp_loss": 0.3077423636526834, "train_patch_loss": 0.9958474041438468, "train_token_loss": 0.7099756244815689, "train_fea_loss": 0.8472568294439885, "epoch": 104, "n_parameters": 144845568}
+{"train_lr": 5.8031985874119795e-05, "train_loss": 2.692192362849232, "train_task_loss": 2.5543478446916565, "train_bpp_loss": 0.30632116043788316, "train_patch_loss": 0.9955702859291927, "train_token_loss": 0.7111749301689099, "train_fea_loss": 0.8476026195617352, "epoch": 105, "n_parameters": 144845568}
+{"train_lr": 5.722897134394433e-05, "train_loss": 2.6670049736075265, "train_task_loss": 2.5291183430889097, "train_bpp_loss": 0.30641474233207694, "train_patch_loss": 0.9873280670623378, "train_token_loss": 0.7049247939295942, "train_fea_loss": 0.8368654731232271, "epoch": 106, "n_parameters": 144845568}
+{"train_lr": 5.642694894242339e-05, "train_loss": 2.6757864850066264, "train_task_loss": 2.538502600463889, "train_bpp_loss": 0.30507530629348983, "train_patch_loss": 0.9907261180832124, "train_token_loss": 0.7074770695045233, "train_fea_loss": 0.8402994043602467, "epoch": 107, "n_parameters": 144845568}
+{"train_lr": 5.562611655657961e-05, "train_loss": 2.7075679923811977, "train_task_loss": 2.5688214136613645, "train_bpp_loss": 0.30832573917316947, "train_patch_loss": 1.00075649227545, "train_token_loss": 0.7133955942328385, "train_fea_loss": 0.8546693176053447, "epoch": 108, "n_parameters": 144845568}
+{"train_lr": 5.642694894242339e-05, "train_loss": 2.668979725284542, "train_task_loss": 2.5311407839669453, "train_bpp_loss": 0.30630876569577653, "train_patch_loss": 0.9885714473649502, "train_token_loss": 0.7049617424553676, "train_fea_loss": 0.8376075848758363, "epoch": 106, "n_parameters": 144845568}
+{"train_lr": 5.642694894242339e-05, "train_loss": 2.683824887423278, "train_task_loss": 2.5468173437535193, "train_bpp_loss": 0.3044612155992725, "train_patch_loss": 0.9947493373275661, "train_token_loss": 0.708695182786714, "train_fea_loss": 0.8433728142076438, "epoch": 107, "n_parameters": 144845568}
+{"train_lr": 5.562611655657961e-05, "train_loss": 2.6808553397548285, "train_task_loss": 2.544291126874568, "train_bpp_loss": 0.30347603688025465, "train_patch_loss": 0.9940348935472093, "train_token_loss": 0.7075042654525676, "train_fea_loss": 0.8427519599440322, "epoch": 108, "n_parameters": 144845568}
+{"train_lr": 5.482667177983261e-05, "train_loss": 4.911457611171962, "train_task_loss": 4.756159201309049, "train_bpp_loss": 0.3451075883260068, "train_patch_loss": 1.923462125296287, "train_token_loss": 0.9228969771633867, "train_fea_loss": 1.909800109163159, "epoch": 109, "n_parameters": 144845568}
+{"train_lr": 5.402881186319929e-05, "train_loss": 6.119308372142075, "train_task_loss": 6.05821056978451, "train_bpp_loss": 0.13577289702430786, "train_patch_loss": 2.0250163967821666, "train_token_loss": 2.0047141713507886, "train_fea_loss": 2.028479987103269, "epoch": 110, "n_parameters": 144845568}
+{"train_lr": 5.402881186319929e-05, "train_loss": 2.8206423869092974, "train_task_loss": 2.5848962855514146, "train_bpp_loss": 0.3167879025416105, "train_patch_loss": 1.0065730715067767, "train_token_loss": 0.721063517386929, "train_fea_loss": 0.8572596858927815, "epoch": 109, "n_parameters": 144845568}
+{"train_lr": 5.402881186319929e-05, "train_loss": 2.7694720208591264, "train_task_loss": 2.536599607177847, "train_bpp_loss": 0.31375666602561586, "train_patch_loss": 0.9914828701552263, "train_token_loss": 0.7029595081839595, "train_fea_loss": 0.8421572199877467, "epoch": 110, "n_parameters": 144845568}
+{"train_lr": 5.323273366669127e-05, "train_loss": 2.754181052810497, "train_task_loss": 2.52081877488098, "train_bpp_loss": 0.31649657639471607, "train_patch_loss": 0.9853969901527683, "train_token_loss": 0.7020405733832817, "train_fea_loss": 0.8333812024247803, "epoch": 111, "n_parameters": 144845568}
+{"train_lr": 5.24386336106797e-05, "train_loss": 2.778903913664089, "train_task_loss": 2.54480576113188, "train_bpp_loss": 0.31611214793115244, "train_patch_loss": 0.9951485451065201, "train_token_loss": 0.7078502580961235, "train_fea_loss": 0.8418069494334306, "epoch": 112, "n_parameters": 144845568}
+{"train_lr": 5.1646707627478925e-05, "train_loss": 2.8131009751854896, "train_task_loss": 2.5725743175803615, "train_bpp_loss": 0.32791142306922616, "train_patch_loss": 1.0041894356025567, "train_token_loss": 0.7132421619497793, "train_fea_loss": 0.8551427105003946, "epoch": 113, "n_parameters": 144845568}
+{"train_lr": 5.0857151112976574e-05, "train_loss": 2.7399135867147137, "train_task_loss": 2.5091776530072987, "train_bpp_loss": 0.31148358721875674, "train_patch_loss": 0.9819968594774175, "train_token_loss": 0.6978094259017913, "train_fea_loss": 0.8293713586163713, "epoch": 114, "n_parameters": 144845568}
+{"train_lr": 5.007015887842505e-05, "train_loss": 2.8684129846825015, "train_task_loss": 2.616686708519427, "train_bpp_loss": 0.3489265601872076, "train_patch_loss": 1.0203611283837783, "train_token_loss": 0.7224983579727063, "train_fea_loss": 0.8738272108821025, "epoch": 115, "n_parameters": 144845568}
+{"train_lr": 4.928592510238729e-05, "train_loss": 2.739907768231144, "train_task_loss": 2.50749061155698, "train_bpp_loss": 0.3151009082838857, "train_patch_loss": 0.9832527024260492, "train_token_loss": 0.695053851908649, "train_fea_loss": 0.8291840486706411, "epoch": 116, "n_parameters": 144845568}
+{"train_lr": 4.850464328279906e-05, "train_loss": 2.7357478904620494, "train_task_loss": 2.5043268277079798, "train_bpp_loss": 0.31352339563690285, "train_patch_loss": 0.9800204708680904, "train_token_loss": 0.6976143020144898, "train_fea_loss": 0.8266920464814639, "epoch": 117, "n_parameters": 144845568}
+{"train_lr": 4.7726506189276635e-05, "train_loss": 2.755742895097541, "train_task_loss": 2.5216859226211086, "train_bpp_loss": 0.31784185127531595, "train_patch_loss": 0.9866870453634815, "train_token_loss": 0.701121356496833, "train_fea_loss": 0.833877512173419, "epoch": 118, "n_parameters": 144845568}
+{"train_lr": 4.69517058154867e-05, "train_loss": 2.7318975441247155, "train_task_loss": 2.5014077417141527, "train_bpp_loss": 0.31166150058537045, "train_patch_loss": 0.9790000828715573, "train_token_loss": 0.6965634005313976, "train_fea_loss": 0.8258442498243256, "epoch": 119, "n_parameters": 144845568}
+{"train_lr": 4.6180433331847694e-05, "train_loss": 2.7518172350307877, "train_task_loss": 2.5213089466577383, "train_bpp_loss": 0.31003568373504375, "train_patch_loss": 0.9861942503458376, "train_token_loss": 0.7014625150290468, "train_fea_loss": 0.8336521727021793, "epoch": 120, "n_parameters": 144845568}
+{"train_lr": 4.541287903828179e-05, "train_loss": 2.7335522819390827, "train_task_loss": 2.50380109216598, "train_bpp_loss": 0.3096768988276289, "train_patch_loss": 0.9810024427442099, "train_token_loss": 0.6958675646743507, "train_fea_loss": 0.8269310759501063, "epoch": 121, "n_parameters": 144845568}
+{"train_lr": 4.4649232317341524e-05, "train_loss": 2.733630260037218, "train_task_loss": 2.503839423107229, "train_bpp_loss": 0.309786735956805, "train_patch_loss": 0.9804957953708278, "train_token_loss": 0.6961014977024256, "train_fea_loss": 0.8272421213170643, "epoch": 122, "n_parameters": 144845568}
+{"train_lr": 4.3889681587425266e-05, "train_loss": 2.8017713390469408, "train_task_loss": 2.557888786087362, "train_bpp_loss": 0.3367434704309995, "train_patch_loss": 0.9978989840482922, "train_token_loss": 0.7109271357421109, "train_fea_loss": 0.8490626564633539, "epoch": 123, "n_parameters": 144845568}
+{"train_lr": 4.313441425631543e-05, "train_loss": 2.7663599788803133, "train_task_loss": 2.515742145719931, "train_bpp_loss": 0.35449836285320296, "train_patch_loss": 0.9864850308973905, "train_token_loss": 0.6938689482427365, "train_fea_loss": 0.8353881579388281, "epoch": 124, "n_parameters": 144845568}
+{"train_lr": 4.238361667491207e-05, "train_loss": 2.7239683468153864, "train_task_loss": 2.4941051970068497, "train_bpp_loss": 0.31064632278029547, "train_patch_loss": 0.9776195450965092, "train_token_loss": 0.6926572750338155, "train_fea_loss": 0.8238283673957955, "epoch": 125, "n_parameters": 144845568}
+{"train_lr": 4.1637474091286196e-05, "train_loss": 2.731808961188193, "train_task_loss": 2.502355479927872, "train_bpp_loss": 0.30911144960618064, "train_patch_loss": 0.9799622440648129, "train_token_loss": 0.6952871139065253, "train_fea_loss": 0.827106113864971, "epoch": 126, "n_parameters": 144845568}
+{"train_lr": 4.089617060496659e-05, "train_loss": 2.7330640450530916, "train_task_loss": 2.5038258624519947, "train_bpp_loss": 0.3084725750156855, "train_patch_loss": 0.9800150552474195, "train_token_loss": 0.6953135876708644, "train_fea_loss": 0.8284972119714609, "epoch": 127, "n_parameters": 144845568}
+{"train_lr": 4.015988912148501e-05, "train_loss": 2.7349733328558417, "train_task_loss": 2.4990430673856814, "train_bpp_loss": 0.3236985370981322, "train_patch_loss": 0.9784446784811054, "train_token_loss": 0.6937227977542497, "train_fea_loss": 0.8268755828197911, "epoch": 128, "n_parameters": 144845568}
+{"train_lr": 3.942881130728865e-05, "train_loss": 2.7178904607856302, "train_task_loss": 2.4889944930085175, "train_bpp_loss": 0.3089114015997885, "train_patch_loss": 0.9757260797004834, "train_token_loss": 0.6912759416524491, "train_fea_loss": 0.821992463240181, "epoch": 129, "n_parameters": 144845568}
+{"train_lr": 3.870311754488397e-05, "train_loss": 2.7325405643301472, "train_task_loss": 2.5016696595584604, "train_bpp_loss": 0.31216120841294276, "train_patch_loss": 0.9799974317856722, "train_token_loss": 0.6937006032407034, "train_fea_loss": 0.8279716155806677, "epoch": 130, "n_parameters": 144845568}
+{"train_lr": 3.798298688834852e-05, "train_loss": 2.7105408391917494, "train_task_loss": 2.4821574514736233, "train_bpp_loss": 0.30818621284714776, "train_patch_loss": 0.9738334429783108, "train_token_loss": 0.6881637847457501, "train_fea_loss": 0.820160214387202, "epoch": 131, "n_parameters": 144845568}
+{"train_lr": 3.726859701914403e-05, "train_loss": 2.7591426904252967, "train_task_loss": 2.5228572249930683, "train_bpp_loss": 0.3226083974237833, "train_patch_loss": 0.9864406956994026, "train_token_loss": 0.7006216509741165, "train_fea_loss": 0.8357948688053184, "epoch": 132, "n_parameters": 144845568}
+{"train_lr": 3.656012420228689e-05, "train_loss": 2.68967245043003, "train_task_loss": 2.460287703357512, "train_bpp_loss": 0.3120911338152163, "train_patch_loss": 0.9669189435883249, "train_token_loss": 0.6814111765313278, "train_fea_loss": 0.81195757473982, "epoch": 133, "n_parameters": 144845568}
+{"train_lr": 3.5857743242838835e-05, "train_loss": 2.7014216272432883, "train_task_loss": 2.472743094181843, "train_bpp_loss": 0.30973817068163295, "train_patch_loss": 0.9706215003429414, "train_token_loss": 0.6868141930465468, "train_fea_loss": 0.8153073933484744, "epoch": 134, "n_parameters": 144845568}
+{"train_lr": 3.516162744279572e-05, "train_loss": 2.705487959069028, "train_task_loss": 2.4765327244067934, "train_bpp_loss": 0.3098553310503527, "train_patch_loss": 0.9728047358189269, "train_token_loss": 0.68612423896057, "train_fea_loss": 0.8176037415066998, "epoch": 135, "n_parameters": 144845568}
+{"train_lr": 3.447194855830639e-05, "train_loss": 2.7113154404383484, "train_task_loss": 2.478181616079321, "train_bpp_loss": 0.31905700233530443, "train_patch_loss": 0.9727772305422514, "train_token_loss": 0.68701637881234, "train_fea_loss": 0.818387999384702, "epoch": 136, "n_parameters": 144845568}
+{"train_lr": 3.378887675732868e-05, "train_loss": 2.6797968578620925, "train_task_loss": 2.451133315422409, "train_bpp_loss": 0.3113255005260612, "train_patch_loss": 0.9640113939207307, "train_token_loss": 0.6797902039285723, "train_fea_loss": 0.8073317096557477, "epoch": 137, "n_parameters": 144845568}
+{"train_lr": 3.311258057759679e-05, "train_loss": 2.694638561546731, "train_task_loss": 2.4661759902551164, "train_bpp_loss": 0.3096901442793062, "train_patch_loss": 0.9691699242545403, "train_token_loss": 0.6841340864242481, "train_fea_loss": 0.8128719716969368, "epoch": 138, "n_parameters": 144845568}
+{"train_lr": 3.244322688507758e-05, "train_loss": 2.697916217782586, "train_task_loss": 2.4694298465224764, "train_bpp_loss": 0.3092930535785854, "train_patch_loss": 0.9713080364482687, "train_token_loss": 0.6833381301698853, "train_fea_loss": 0.8147836721902021, "epoch": 139, "n_parameters": 144845568}
+{"train_lr": 3.1780980832784374e-05, "train_loss": 2.697745393452456, "train_task_loss": 2.465554346515835, "train_bpp_loss": 0.3177864467847994, "train_patch_loss": 0.9703194021450565, "train_token_loss": 0.6818096557480219, "train_fea_loss": 0.8134252810515732, "epoch": 140, "n_parameters": 144845568}
+{"train_lr": 3.112600582001298e-05, "train_loss": 2.692317468710512, "train_task_loss": 2.4639319612039365, "train_bpp_loss": 0.3095171678927448, "train_patch_loss": 0.9694184379415404, "train_token_loss": 0.6818745398862684, "train_fea_loss": 0.8126389763923251, "epoch": 141, "n_parameters": 144845568}
+{"train_lr": 3.047846345205177e-05, "train_loss": 2.7029235281175277, "train_task_loss": 2.47447449285135, "train_bpp_loss": 0.30879852916527206, "train_patch_loss": 0.9717401454551257, "train_token_loss": 0.6846787373957445, "train_fea_loss": 0.8180556024426965, "epoch": 142, "n_parameters": 144845568}
+{"train_lr": 2.9838513500286588e-05, "train_loss": 2.6954629459469723, "train_task_loss": 2.461703634662308, "train_bpp_loss": 0.3215753604463643, "train_patch_loss": 0.9685670451911019, "train_token_loss": 0.6806938354801432, "train_fea_loss": 0.8124427453947546, "epoch": 143, "n_parameters": 144845568}
+{"train_lr": 2.920631386279756e-05, "train_loss": 2.6789074894978846, "train_task_loss": 2.4506222840836296, "train_bpp_loss": 0.3103648301312249, "train_patch_loss": 0.9643291673049045, "train_token_loss": 0.6781998446358622, "train_fea_loss": 0.8080932640043987, "epoch": 144, "n_parameters": 144845568}
+{"train_lr": 2.8582020525382766e-05, "train_loss": 2.698681264338519, "train_task_loss": 2.465655054361057, "train_bpp_loss": 0.3195096456470321, "train_patch_loss": 0.9703498526188217, "train_token_loss": 0.6807158637626601, "train_fea_loss": 0.8145893300555164, "epoch": 145, "n_parameters": 144845568}
+{"train_lr": 2.7965787523079142e-05, "train_loss": 2.6876495937845832, "train_task_loss": 2.4591161769142538, "train_bpp_loss": 0.3100968778760974, "train_patch_loss": 0.9675739941743317, "train_token_loss": 0.6793179688975215, "train_fea_loss": 0.8122242059786435, "epoch": 146, "n_parameters": 144845568}
+{"train_lr": 2.7357766902161244e-05, "train_loss": 2.686950199314945, "train_task_loss": 2.4579590937520246, "train_bpp_loss": 0.31120623359166316, "train_patch_loss": 0.9676960334248788, "train_token_loss": 0.6789913053403227, "train_fea_loss": 0.8112717479337855, "epoch": 147, "n_parameters": 144845568}
+{"train_lr": 2.616696082115359e-05, "train_loss": 2.6700563189937627, "train_task_loss": 2.441784028437355, "train_bpp_loss": 0.31092876969329675, "train_patch_loss": 0.9618561120337541, "train_token_loss": 0.6746953771956604, "train_fea_loss": 0.8052325316769577, "epoch": 148, "n_parameters": 144845568}
+{"train_lr": 2.616696082115359e-05, "train_loss": 2.6674406088138225, "train_task_loss": 2.4389501352229424, "train_bpp_loss": 0.3116684112110459, "train_patch_loss": 0.9607490962326563, "train_token_loss": 0.6741545937120093, "train_fea_loss": 0.8040464378075955, "epoch": 149, "n_parameters": 144845568}
+{"train_lr": 2.558446917464184e-05, "train_loss": 2.6966678162129947, "train_task_loss": 2.4605731371191624, "train_bpp_loss": 0.32649256723287473, "train_patch_loss": 0.9711854177922439, "train_token_loss": 0.6771108717668792, "train_fea_loss": 0.81227684003129, "epoch": 150, "n_parameters": 144845568}
+{"train_lr": 2.5010777464192224e-05, "train_loss": 2.649356059051103, "train_task_loss": 2.4205873620649463, "train_bpp_loss": 0.3137229763348844, "train_patch_loss": 0.9557904277322925, "train_token_loss": 0.6687184504994885, "train_fea_loss": 0.7960784757974372, "epoch": 151, "n_parameters": 144845568}
+{"train_lr": 2.444602723963776e-05, "train_loss": 2.6504098884582663, "train_task_loss": 2.421797823491428, "train_bpp_loss": 0.3133131659475805, "train_patch_loss": 0.9554594567903191, "train_token_loss": 0.6693732499343743, "train_fea_loss": 0.7969651087689743, "epoch": 152, "n_parameters": 144845568}
+{"train_lr": 2.3343906382349e-05, "train_loss": 2.7114228718429447, "train_task_loss": 2.458482252813572, "train_bpp_loss": 0.34088637049106674, "train_patch_loss": 0.9697519682405843, "train_token_loss": 0.6721772563775523, "train_fea_loss": 0.8165530218369812, "epoch": 153, "n_parameters": 144845568}
+{"train_lr": 2.3343906382349e-05, "train_loss": 2.650369150729345, "train_task_loss": 2.4105861190197277, "train_bpp_loss": 0.31759558117231557, "train_patch_loss": 0.9528444103351671, "train_token_loss": 0.6638412714524896, "train_fea_loss": 0.7939004297111388, "epoch": 154, "n_parameters": 144845568}
+{"train_lr": 2.280680768143689e-05, "train_loss": 2.658987253216459, "train_task_loss": 2.421917999400009, "train_bpp_loss": 0.31114791267567116, "train_patch_loss": 0.9561528907575124, "train_token_loss": 0.6675511025901821, "train_fea_loss": 0.7982139991355135, "epoch": 155, "n_parameters": 144845568}
+{"train_lr": 2.2279194262997928e-05, "train_loss": 2.695716130978269, "train_task_loss": 2.4542629720031215, "train_bpp_loss": 0.31765079001362306, "train_patch_loss": 0.9665458119615055, "train_token_loss": 0.6746455496869428, "train_fea_loss": 0.8130716029219598, "epoch": 156, "n_parameters": 144845568}
+{"train_lr": 2.1761196307742086e-05, "train_loss": 2.6704102983202436, "train_task_loss": 2.4334238473442102, "train_bpp_loss": 0.30996573550349693, "train_patch_loss": 0.9596143975134769, "train_token_loss": 0.6693641086682963, "train_fea_loss": 0.8044453353511629, "epoch": 157, "n_parameters": 144845568}
+{"train_lr": 2.1252941623912577e-05, "train_loss": 2.6702361341735585, "train_task_loss": 2.4332374961917207, "train_bpp_loss": 0.3099473466777312, "train_patch_loss": 0.9602035106019043, "train_token_loss": 0.6687574936278003, "train_fea_loss": 0.8042764851122165, "epoch": 158, "n_parameters": 144845568}
+{"train_lr": 2.0754555615745688e-05, "train_loss": 2.6676239087480864, "train_task_loss": 2.430506248768571, "train_bpp_loss": 0.31043111976030635, "train_patch_loss": 0.9591814159645701, "train_token_loss": 0.6682900374808864, "train_fea_loss": 0.8030347887748223, "epoch": 159, "n_parameters": 144845568}
+{"train_lr": 2.0266161252534863e-05, "train_loss": 2.682328796143726, "train_task_loss": 2.4449838221484095, "train_bpp_loss": 0.3097276722284244, "train_patch_loss": 0.9639880869501858, "train_token_loss": 0.6714682123563487, "train_fea_loss": 0.8095275162444483, "epoch": 160, "n_parameters": 144845568}
+{"train_lr": 1.9787879038283694e-05, "train_loss": 2.794825274291799, "train_task_loss": 2.5297502170381643, "train_bpp_loss": 0.3605138300236859, "train_patch_loss": 0.9950330435751761, "train_token_loss": 0.6891800253216526, "train_fea_loss": 0.8455371403780808, "epoch": 161, "n_parameters": 144845568}
+{"train_lr": 1.9319826981968032e-05, "train_loss": 2.6933876661052234, "train_task_loss": 2.4523086957990836, "train_bpp_loss": 0.3166517836579995, "train_patch_loss": 0.9668755698480087, "train_token_loss": 0.6705846348107397, "train_fea_loss": 0.8148484852687596, "epoch": 162, "n_parameters": 144845568}
+{"train_lr": 1.8862120568428674e-05, "train_loss": 2.7029373263426537, "train_task_loss": 2.4646165441873547, "train_bpp_loss": 0.31019026768110175, "train_patch_loss": 0.9700850130145927, "train_token_loss": 0.6760256816487875, "train_fea_loss": 0.8185058429220812, "epoch": 163, "n_parameters": 144845568}
+{"train_lr": 1.8414872729877464e-05, "train_loss": 2.7035999098252192, "train_task_loss": 2.4655938773311012, "train_bpp_loss": 0.30946591779214444, "train_patch_loss": 0.9705398866317148, "train_token_loss": 0.6763441974645574, "train_fea_loss": 0.8187097842231107, "epoch": 164, "n_parameters": 144845568}
+{"train_lr": 1.9319826981970397e-05, "train_loss": 2.740814351822904, "train_task_loss": 2.492047573066801, "train_bpp_loss": 0.329818556624592, "train_patch_loss": 0.9823994228093744, "train_token_loss": 0.6829682969436836, "train_fea_loss": 0.8266798452512144, "epoch": 161, "n_parameters": 144845568}
+{"train_lr": 1.9319826981970397e-05, "train_loss": 2.7282523382946455, "train_task_loss": 2.48745193989943, "train_bpp_loss": 0.3139055231898773, "train_patch_loss": 0.9767762538854131, "train_token_loss": 0.6849351603363594, "train_fea_loss": 0.8257405170858156, "epoch": 162, "n_parameters": 144845568}
+{"train_lr": 1.8862120568426702e-05, "train_loss": 2.7373931265259674, "train_task_loss": 2.493901286015122, "train_bpp_loss": 0.3186548652218954, "train_patch_loss": 0.9806770237019773, "train_token_loss": 0.6831493455675437, "train_fea_loss": 0.8300749089220445, "epoch": 163, "n_parameters": 144845568}
+{"train_lr": 1.84148727298801e-05, "train_loss": 2.7362954941370505, "train_task_loss": 2.4902694111319184, "train_bpp_loss": 0.3243905476437395, "train_patch_loss": 0.9797143076258383, "train_token_loss": 0.6838951857773949, "train_fea_loss": 0.8266599099073622, "epoch": 164, "n_parameters": 144845568}

1_feature_extractor/log/DINOv2_training/log/20240725_084736.log

@@ -0,0 +1,555 @@
+2024-07-25 08:47:36,358 [INFO ]  Logging file is /data0/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/log//20240725_084736.log
+2024-07-25 08:47:36,359 [INFO ]  job dir: /data0/qiyp/Proteus-pytorch/pretrain
+2024-07-25 08:47:36,359 [INFO ]  Namespace(batch_size=48,
+epochs=200,
+bce_loss=False,
+unscale_lr=False,
+model='models_proteus_dinov2',
+target_model='vit_base',
+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.0002,
+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,
+global_crops_size=224,
+patch_size=14,
+mask_ratio=[0.5],
+mask_probability=0.5,
+mask_first_n=True,
+clone_batch=1,
+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='vit_large',
+teacher_path='',
+distillation_type='none',
+distillation_alpha=0.5,
+distillation_tau=1.0,
+lambda_token=1.0,
+lambda_fea=1.0,
+lambda_patch=1.0,
+cosub=False,
+finetune='/data0/qiyp/Proteus-pytorch/pretrain/ckpt/proteus_vitb_backbone.pth',
+attn_only=False,
+weight_inherit='',
+data_path='/data1/datasets/imagenet_fold',
+data_set='IMNET',
+inat_category='name',
+output_dir='log/DINOv2_training',
+log_dir='/data0/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/log/',
+device='cuda',
+seed=0,
+resume='',
+start_epoch=0,
+eval=False,
+eval_crop_ratio=0.875,
+dist_eval=False,
+num_workers=10,
+pin_mem=True,
+distributed=True,
+world_size=6,
+dist_url='env://',
+rank=0,
+gpu=0,
+dist_backend='nccl')
+
+2024-07-25 08:47:38,916 [INFO ]  Dataset ImageFolder
+    Number of datapoints: 1281167
+    Root location: /data1/datasets/imagenet_fold/train
+    StandardTransform
+Transform: Compose(
+               RandomResizedCropAndInterpolation(size=(224, 224), scale=(0.08, 1.0), ratio=(0.75, 1.3333), interpolation=bicubic)
+               RandomHorizontalFlip(p=0.5)
+               RandAugment(n=2, ops=
+           	AugmentOp(name=AutoContrast, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=Equalize, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=Invert, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=Rotate, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=PosterizeIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=SolarizeIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=SolarizeAdd, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=ColorIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=ContrastIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=BrightnessIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=SharpnessIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=ShearX, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=ShearY, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=TranslateXRel, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=TranslateYRel, p=0.5, m=9, mstd=0.5))
+               ToTensor()
+               Normalize(mean=tensor([0.4850, 0.4560, 0.4060]), std=tensor([0.2290, 0.2240, 0.2250]))
+               RandomErasing(p=0.25, mode=pixel, count=(1, 1))
+           )
+2024-07-25 08:47:38,926 [INFO ]  Sampler_train = <samplers.RASampler object at 0x7fc548b7f3a0>
+2024-07-25 08:47:38,928 [INFO ]  Created a temporary directory at /tmp/tmpe0qnk5ic
+2024-07-25 08:47:38,929 [INFO ]  Writing /tmp/tmpe0qnk5ic/_remote_module_non_scriptable.py
+2024-07-25 08:47:40,640 [INFO ]  using MLP layer as FFN
+2024-07-25 08:49:51,514 [INFO ]  using MLP layer as FFN
+2024-07-25 08:49:57,349 [INFO ]  Model = MetaArch(
+  (student): ModuleDict(
+    (backbone): DinoVisionTransformer(
+      (patch_embed): PatchEmbed(
+        (proj): Conv2d(3, 768, kernel_size=(14, 14), stride=(14, 14))
+        (norm): Identity()
+      )
+      (blocks): ModuleList(
+        (0-11): 12 x Block(
+          (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
+          (attn): MemEffAttention(
+            (qkv): Linear(in_features=768, out_features=2304, bias=True)
+            (attn_drop): Dropout(p=0.0, inplace=False)
+            (proj): Linear(in_features=768, out_features=768, bias=True)
+            (proj_drop): Dropout(p=0.0, inplace=False)
+          )
+          (ls1): LayerScale()
+          (drop_path1): Identity()
+          (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
+          (mlp): Mlp(
+            (fc1): Linear(in_features=768, out_features=3072, bias=True)
+            (act): GELU(approximate='none')
+            (fc2): Linear(in_features=3072, out_features=768, bias=True)
+            (drop): Dropout(p=0.0, inplace=False)
+          )
+          (ls2): LayerScale()
+          (drop_path2): Identity()
+        )
+      )
+      (norm): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
+      (head): Identity()
+    )
+  )
+  (teacher): ModuleDict(
+    (backbone): DinoVisionTransformer(
+      (patch_embed): PatchEmbed(
+        (proj): Conv2d(3, 1024, kernel_size=(14, 14), stride=(14, 14))
+        (norm): Identity()
+      )
+      (blocks): ModuleList(
+        (0-23): 24 x NestedTensorBlock(
+          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
+          (attn): MemEffAttention(
+            (qkv): Linear(in_features=1024, out_features=3072, bias=True)
+            (attn_drop): Dropout(p=0.0, inplace=False)
+            (proj): Linear(in_features=1024, out_features=1024, bias=True)
+            (proj_drop): Dropout(p=0.0, inplace=False)
+          )
+          (ls1): LayerScale()
+          (drop_path1): Identity()
+          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
+          (mlp): Mlp(
+            (fc1): Linear(in_features=1024, out_features=4096, bias=True)
+            (act): GELU(approximate='none')
+            (fc2): Linear(in_features=4096, out_features=1024, bias=True)
+            (drop): Dropout(p=0.0, inplace=False)
+          )
+          (ls2): LayerScale()
+          (drop_path2): Identity()
+        )
+      )
+      (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
+      (head): Identity()
+    )
+  )
+  (patch_head): Sequential(
+    (0): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (1): Linear(in_features=768, out_features=1024, bias=True)
+  )
+  (token_head): Sequential(
+    (0): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (1): Linear(in_features=768, out_features=1024, bias=True)
+  )
+  (fea_head): Sequential(
+    (0): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (1): Linear(in_features=768, out_features=1024, bias=True)
+  )
+  (soft_criterion): MSELoss()
+  (info_bottleneck): IF_Module(
+    (encoder_blocks): ModuleList(
+      (0-3): 4 x Block(
+        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+        (attn): Attention(
+          (qkv): Linear(in_features=768, out_features=2304, bias=True)
+          (q_norm): Identity()
+          (k_norm): Identity()
+          (attn_drop): Dropout(p=0.0, inplace=False)
+          (proj): Linear(in_features=768, out_features=768, bias=True)
+          (proj_drop): Dropout(p=0.0, inplace=False)
+        )
+        (ls1): Identity()
+        (drop_path1): Identity()
+        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+        (mlp): Mlp(
+          (fc1): Linear(in_features=768, out_features=3072, bias=True)
+          (act): GELU(approximate='none')
+          (drop1): Dropout(p=0.0, inplace=False)
+          (norm): Identity()
+          (fc2): Linear(in_features=3072, out_features=768, bias=True)
+          (drop2): Dropout(p=0.0, inplace=False)
+        )
+        (ls2): Identity()
+        (drop_path2): Identity()
+      )
+    )
+    (encoder_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (decoder_blocks): ModuleList(
+      (0-3): 4 x Block(
+        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+        (attn): Attention(
+          (qkv): Linear(in_features=768, out_features=2304, bias=True)
+          (q_norm): Identity()
+          (k_norm): Identity()
+          (attn_drop): Dropout(p=0.0, inplace=False)
+          (proj): Linear(in_features=768, out_features=768, bias=True)
+          (proj_drop): Dropout(p=0.0, inplace=False)
+        )
+        (ls1): Identity()
+        (drop_path1): Identity()
+        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+        (mlp): Mlp(
+          (fc1): Linear(in_features=768, out_features=3072, bias=True)
+          (act): GELU(approximate='none')
+          (drop1): Dropout(p=0.0, inplace=False)
+          (norm): Identity()
+          (fc2): Linear(in_features=3072, out_features=768, bias=True)
+          (drop2): Dropout(p=0.0, inplace=False)
+        )
+        (ls2): Identity()
+        (drop_path2): Identity()
+      )
+    )
+    (decoder_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (entropy_bottleneck): EntropyBottleneck(
+      (likelihood_lower_bound): LowerBound()
+    )
+  )
+)
+2024-07-25 08:49:59,702 [INFO ]  Finetuning from /data0/qiyp/Proteus-pytorch/pretrain/ckpt/proteus_vitb_backbone.pth
+2024-07-25 08:49:59,703 [INFO ]  missing_keys: ['patch_head.0.weight', 'patch_head.0.bias', 'patch_head.1.weight', 'patch_head.1.bias', 'info_bottleneck.encoder_blocks.0.norm1.weight', 'info_bottleneck.encoder_blocks.0.norm1.bias', 'info_bottleneck.encoder_blocks.0.attn.qkv.weight', 'info_bottleneck.encoder_blocks.0.attn.qkv.bias', 'info_bottleneck.encoder_blocks.0.attn.proj.weight', 'info_bottleneck.encoder_blocks.0.attn.proj.bias', 'info_bottleneck.encoder_blocks.0.norm2.weight', 'info_bottleneck.encoder_blocks.0.norm2.bias', 'info_bottleneck.encoder_blocks.0.mlp.fc1.weight', 'info_bottleneck.encoder_blocks.0.mlp.fc1.bias', 'info_bottleneck.encoder_blocks.0.mlp.fc2.weight', 'info_bottleneck.encoder_blocks.0.mlp.fc2.bias', 'info_bottleneck.encoder_blocks.1.norm1.weight', 'info_bottleneck.encoder_blocks.1.norm1.bias', 'info_bottleneck.encoder_blocks.1.attn.qkv.weight', 'info_bottleneck.encoder_blocks.1.attn.qkv.bias', 'info_bottleneck.encoder_blocks.1.attn.proj.weight', 'info_bottleneck.encoder_blocks.1.attn.proj.bias', 'info_bottleneck.encoder_blocks.1.norm2.weight', 'info_bottleneck.encoder_blocks.1.norm2.bias', 'info_bottleneck.encoder_blocks.1.mlp.fc1.weight', 'info_bottleneck.encoder_blocks.1.mlp.fc1.bias', 'info_bottleneck.encoder_blocks.1.mlp.fc2.weight', 'info_bottleneck.encoder_blocks.1.mlp.fc2.bias', 'info_bottleneck.encoder_blocks.2.norm1.weight', 'info_bottleneck.encoder_blocks.2.norm1.bias', 'info_bottleneck.encoder_blocks.2.attn.qkv.weight', 'info_bottleneck.encoder_blocks.2.attn.qkv.bias', 'info_bottleneck.encoder_blocks.2.attn.proj.weight', 'info_bottleneck.encoder_blocks.2.attn.proj.bias', 'info_bottleneck.encoder_blocks.2.norm2.weight', 'info_bottleneck.encoder_blocks.2.norm2.bias', 'info_bottleneck.encoder_blocks.2.mlp.fc1.weight', 'info_bottleneck.encoder_blocks.2.mlp.fc1.bias', 'info_bottleneck.encoder_blocks.2.mlp.fc2.weight', 'info_bottleneck.encoder_blocks.2.mlp.fc2.bias', 'info_bottleneck.encoder_blocks.3.norm1.weight', 'info_bottleneck.encoder_blocks.3.norm1.bias', 'info_bottleneck.encoder_blocks.3.attn.qkv.weight', 'info_bottleneck.encoder_blocks.3.attn.qkv.bias', 'info_bottleneck.encoder_blocks.3.attn.proj.weight', 'info_bottleneck.encoder_blocks.3.attn.proj.bias', 'info_bottleneck.encoder_blocks.3.norm2.weight', 'info_bottleneck.encoder_blocks.3.norm2.bias', 'info_bottleneck.encoder_blocks.3.mlp.fc1.weight', 'info_bottleneck.encoder_blocks.3.mlp.fc1.bias', 'info_bottleneck.encoder_blocks.3.mlp.fc2.weight', 'info_bottleneck.encoder_blocks.3.mlp.fc2.bias', 'info_bottleneck.encoder_norm.weight', 'info_bottleneck.encoder_norm.bias', 'info_bottleneck.decoder_blocks.0.norm1.weight', 'info_bottleneck.decoder_blocks.0.norm1.bias', 'info_bottleneck.decoder_blocks.0.attn.qkv.weight', 'info_bottleneck.decoder_blocks.0.attn.qkv.bias', 'info_bottleneck.decoder_blocks.0.attn.proj.weight', 'info_bottleneck.decoder_blocks.0.attn.proj.bias', 'info_bottleneck.decoder_blocks.0.norm2.weight', 'info_bottleneck.decoder_blocks.0.norm2.bias', 'info_bottleneck.decoder_blocks.0.mlp.fc1.weight', 'info_bottleneck.decoder_blocks.0.mlp.fc1.bias', 'info_bottleneck.decoder_blocks.0.mlp.fc2.weight', 'info_bottleneck.decoder_blocks.0.mlp.fc2.bias', 'info_bottleneck.decoder_blocks.1.norm1.weight', 'info_bottleneck.decoder_blocks.1.norm1.bias', 'info_bottleneck.decoder_blocks.1.attn.qkv.weight', 'info_bottleneck.decoder_blocks.1.attn.qkv.bias', 'info_bottleneck.decoder_blocks.1.attn.proj.weight', 'info_bottleneck.decoder_blocks.1.attn.proj.bias', 'info_bottleneck.decoder_blocks.1.norm2.weight', 'info_bottleneck.decoder_blocks.1.norm2.bias', 'info_bottleneck.decoder_blocks.1.mlp.fc1.weight', 'info_bottleneck.decoder_blocks.1.mlp.fc1.bias', 'info_bottleneck.decoder_blocks.1.mlp.fc2.weight', 'info_bottleneck.decoder_blocks.1.mlp.fc2.bias', 'info_bottleneck.decoder_blocks.2.norm1.weight', 'info_bottleneck.decoder_blocks.2.norm1.bias', 'info_bottleneck.decoder_blocks.2.attn.qkv.weight', 'info_bottleneck.decoder_blocks.2.attn.qkv.bias', 'info_bottleneck.decoder_blocks.2.attn.proj.weight', 'info_bottleneck.decoder_blocks.2.attn.proj.bias', 'info_bottleneck.decoder_blocks.2.norm2.weight', 'info_bottleneck.decoder_blocks.2.norm2.bias', 'info_bottleneck.decoder_blocks.2.mlp.fc1.weight', 'info_bottleneck.decoder_blocks.2.mlp.fc1.bias', 'info_bottleneck.decoder_blocks.2.mlp.fc2.weight', 'info_bottleneck.decoder_blocks.2.mlp.fc2.bias', 'info_bottleneck.decoder_blocks.3.norm1.weight', 'info_bottleneck.decoder_blocks.3.norm1.bias', 'info_bottleneck.decoder_blocks.3.attn.qkv.weight', 'info_bottleneck.decoder_blocks.3.attn.qkv.bias', 'info_bottleneck.decoder_blocks.3.attn.proj.weight', 'info_bottleneck.decoder_blocks.3.attn.proj.bias', 'info_bottleneck.decoder_blocks.3.norm2.weight', 'info_bottleneck.decoder_blocks.3.norm2.bias', 'info_bottleneck.decoder_blocks.3.mlp.fc1.weight', 'info_bottleneck.decoder_blocks.3.mlp.fc1.bias', 'info_bottleneck.decoder_blocks.3.mlp.fc2.weight', 'info_bottleneck.decoder_blocks.3.mlp.fc2.bias', 'info_bottleneck.decoder_norm.weight', 'info_bottleneck.decoder_norm.bias', 'info_bottleneck.entropy_bottleneck._matrix0', 'info_bottleneck.entropy_bottleneck._bias0', 'info_bottleneck.entropy_bottleneck._factor0', 'info_bottleneck.entropy_bottleneck._matrix1', 'info_bottleneck.entropy_bottleneck._bias1', 'info_bottleneck.entropy_bottleneck._factor1', 'info_bottleneck.entropy_bottleneck._matrix2', 'info_bottleneck.entropy_bottleneck._bias2', 'info_bottleneck.entropy_bottleneck._factor2', 'info_bottleneck.entropy_bottleneck._matrix3', 'info_bottleneck.entropy_bottleneck._bias3', 'info_bottleneck.entropy_bottleneck._factor3', 'info_bottleneck.entropy_bottleneck._matrix4', 'info_bottleneck.entropy_bottleneck._bias4', 'info_bottleneck.entropy_bottleneck.quantiles', 'info_bottleneck.entropy_bottleneck._offset', 'info_bottleneck.entropy_bottleneck._quantized_cdf', 'info_bottleneck.entropy_bottleneck._cdf_length', 'info_bottleneck.entropy_bottleneck.target', 'info_bottleneck.entropy_bottleneck.likelihood_lower_bound.bound']
+2024-07-25 08:49:59,703 [INFO ]  unexpected_keys: ['ibot_head.0.weight', 'ibot_head.0.bias', 'ibot_head.1.weight', 'ibot_head.1.bias']
+2024-07-25 08:50:00,636 [INFO ]  number of params: 144845568
+2024-07-25 08:50:00,636 [INFO ]  base lr = 0.0002
+2024-07-25 08:50:00,636 [INFO ]  actural lr = 0.00011250000000000001
+2024-07-25 08:50:00,639 [INFO ]  Start training for 200 epochs
+2024-07-25 08:50:00,642 [INFO ]  Parameters to be updated: 
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.0.norm2.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.2.norm2.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.1.attn.proj.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.1.attn.qkv.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.1.norm2.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.0.norm2.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.1.mlp.fc2.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.entropy_bottleneck._bias0
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.2.attn.proj.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.1.mlp.fc2.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.0.norm1.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.entropy_bottleneck._matrix0
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.entropy_bottleneck._factor2
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.3.attn.proj.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.1.attn.proj.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.2.norm1.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.1.mlp.fc2.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.entropy_bottleneck._matrix4
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.3.mlp.fc2.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.3.norm1.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.0.mlp.fc2.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.0.norm1.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.2.norm2.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.3.norm1.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.entropy_bottleneck._bias2
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.2.attn.qkv.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.3.attn.qkv.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.1.norm1.bias
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.1.attn.qkv.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.0.norm2.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.decoder_blocks.2.attn.proj.weight
+
+2024-07-25 08:50:00,642 [INFO ]  	module.info_bottleneck.encoder_blocks.2.mlp.fc2.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.3.attn.proj.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.3.mlp.fc2.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.3.norm1.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.0.mlp.fc1.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.entropy_bottleneck._bias1
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.1.norm2.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.0.attn.qkv.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.2.norm2.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.2.attn.qkv.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.1.norm1.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.1.mlp.fc1.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.1.mlp.fc1.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.2.mlp.fc1.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.entropy_bottleneck._bias3
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.2.norm2.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.2.mlp.fc1.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.1.norm2.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.0.norm1.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.0.attn.qkv.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.3.norm2.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.3.mlp.fc2.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.0.mlp.fc1.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.1.attn.proj.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_norm.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_norm.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.3.mlp.fc2.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.3.norm2.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.1.norm1.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.entropy_bottleneck.quantiles
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.entropy_bottleneck._matrix2
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.decoder_blocks.0.attn.proj.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.3.attn.qkv.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.3.mlp.fc1.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.2.attn.proj.weight
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.3.norm2.bias
+
+2024-07-25 08:50:00,643 [INFO ]  	module.info_bottleneck.encoder_blocks.2.mlp.fc2.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.2.mlp.fc2.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.3.attn.proj.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.2.attn.qkv.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.entropy_bottleneck._matrix1
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.3.attn.proj.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.0.attn.qkv.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.1.mlp.fc2.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.0.mlp.fc2.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.1.norm1.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.0.attn.proj.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.2.mlp.fc2.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.2.mlp.fc1.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.1.mlp.fc1.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.0.norm1.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.3.attn.qkv.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.3.norm1.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.2.attn.qkv.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.2.attn.proj.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.2.norm1.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.0.mlp.fc2.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.1.mlp.fc1.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.2.mlp.fc1.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.3.attn.qkv.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.entropy_bottleneck._matrix3
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_norm.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.1.norm2.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.0.mlp.fc2.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.0.mlp.fc1.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.0.attn.proj.bias
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.entropy_bottleneck._factor0
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.1.attn.proj.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.entropy_bottleneck._bias4
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.entropy_bottleneck._factor3
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.encoder_blocks.0.mlp.fc1.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.2.norm1.weight
+
+2024-07-25 08:50:00,644 [INFO ]  	module.info_bottleneck.decoder_blocks.3.norm2.bias
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.decoder_norm.weight
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.decoder_blocks.0.norm2.bias
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.decoder_blocks.1.attn.qkv.bias
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.decoder_blocks.2.norm1.bias
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.decoder_blocks.0.attn.qkv.weight
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.encoder_blocks.3.mlp.fc1.bias
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.decoder_blocks.3.mlp.fc1.weight
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.encoder_blocks.0.attn.proj.bias
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.encoder_blocks.1.attn.qkv.bias
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.decoder_blocks.3.mlp.fc1.bias
+
+2024-07-25 08:50:00,645 [INFO ]  	module.info_bottleneck.entropy_bottleneck._factor1
+
+2024-07-25 08:50:00,645 [INFO ]  
+
+2024-07-25 08:50:06,715 [INFO ]  Epoch: [0]  [   0/4448]  eta: 7:29:47  lr: 0.000001  loss: 8.3333 (8.3333)  task_loss: 5.5091 (5.5091)  bpp_loss: 28.2426 (28.2426)  patch_loss: 2.4117 (2.4117)  token_loss: 0.6234 (0.6234)  fea_loss: 2.4741 (2.4741)  time: 6.0673  data: 2.4582  max mem: 10534
+2024-07-25 08:50:16,045 [INFO ]  Epoch: [0]  [  10/4448]  eta: 1:43:31  lr: 0.000001  loss: 8.3085 (8.3065)  task_loss: 5.4842 (5.4822)  bpp_loss: 28.2425 (28.2426)  patch_loss: 2.4007 (2.4001)  token_loss: 0.6673 (0.6727)  fea_loss: 2.4212 (2.4095)  time: 1.3996  data: 0.2236  max mem: 11190
+2024-07-25 08:50:25,369 [INFO ]  Epoch: [0]  [  20/4448]  eta: 1:26:52  lr: 0.000001  loss: 8.2112 (8.2280)  task_loss: 5.3869 (5.4038)  bpp_loss: 28.2425 (28.2427)  patch_loss: 2.3943 (2.3977)  token_loss: 0.6668 (0.6682)  fea_loss: 2.3140 (2.3378)  time: 0.9326  data: 0.0001  max mem: 11191
+2024-07-25 08:50:34,728 [INFO ]  Epoch: [0]  [  30/4448]  eta: 1:20:56  lr: 0.000001  loss: 8.0968 (8.1642)  task_loss: 5.2725 (5.3399)  bpp_loss: 28.2428 (28.2427)  patch_loss: 2.3875 (2.3914)  token_loss: 0.6681 (0.6704)  fea_loss: 2.2053 (2.2782)  time: 0.9341  data: 0.0001  max mem: 11191
+2024-07-25 08:50:44,122 [INFO ]  Epoch: [0]  [  40/4448]  eta: 1:17:53  lr: 0.000001  loss: 7.9714 (8.1093)  task_loss: 5.1472 (5.2851)  bpp_loss: 28.2426 (28.2427)  patch_loss: 2.3728 (2.3870)  token_loss: 0.6831 (0.6721)  fea_loss: 2.0939 (2.2259)  time: 0.9375  data: 0.0001  max mem: 11191
+2024-07-25 08:50:53,546 [INFO ]  Epoch: [0]  [  50/4448]  eta: 1:16:01  lr: 0.000001  loss: 7.8988 (8.0590)  task_loss: 5.0745 (5.2347)  bpp_loss: 28.2426 (28.2427)  patch_loss: 2.3717 (2.3832)  token_loss: 0.6707 (0.6729)  fea_loss: 2.0184 (2.1786)  time: 0.9408  data: 0.0001  max mem: 11192
+2024-07-25 08:51:02,993 [INFO ]  Epoch: [0]  [  60/4448]  eta: 1:14:44  lr: 0.000001  loss: 7.8305 (8.0147)  task_loss: 5.0062 (5.1904)  bpp_loss: 28.2426 (28.2427)  patch_loss: 2.3574 (2.3787)  token_loss: 0.6676 (0.6735)  fea_loss: 1.9486 (2.1383)  time: 0.9435  data: 0.0001  max mem: 11192
+2024-07-25 08:51:12,484 [INFO ]  Epoch: [0]  [  70/4448]  eta: 1:13:49  lr: 0.000001  loss: 7.7545 (7.9697)  task_loss: 4.9303 (5.1454)  bpp_loss: 28.2426 (28.2427)  patch_loss: 2.3489 (2.3734)  token_loss: 0.6659 (0.6675)  fea_loss: 1.9142 (2.1045)  time: 0.9469  data: 0.0001  max mem: 11192
+2024-07-25 08:51:21,991 [INFO ]  Epoch: [0]  [  80/4448]  eta: 1:13:06  lr: 0.000001  loss: 7.6890 (7.9328)  task_loss: 4.8647 (5.1085)  bpp_loss: 28.2428 (28.2427)  patch_loss: 2.3404 (2.3694)  token_loss: 0.6583 (0.6661)  fea_loss: 1.8710 (2.0731)  time: 0.9498  data: 0.0001  max mem: 11194
+2024-07-25 08:51:31,508 [INFO ]  Epoch: [0]  [  90/4448]  eta: 1:12:30  lr: 0.000001  loss: 7.6696 (7.9028)  task_loss: 4.8453 (5.0785)  bpp_loss: 28.2426 (28.2427)  patch_loss: 2.3398 (2.3655)  token_loss: 0.6613 (0.6665)  fea_loss: 1.8399 (2.0465)  time: 0.9511  data: 0.0001  max mem: 11194
+2024-07-25 08:51:41,113 [INFO ]  Epoch: [0]  [ 100/4448]  eta: 1:12:04  lr: 0.000001  loss: 7.6273 (7.8742)  task_loss: 4.8031 (5.0499)  bpp_loss: 28.2425 (28.2427)  patch_loss: 2.3329 (2.3620)  token_loss: 0.6699 (0.6671)  fea_loss: 1.8089 (2.0209)  time: 0.9560  data: 0.0001  max mem: 11194
+2024-07-25 08:51:50,804 [INFO ]  Epoch: [0]  [ 110/4448]  eta: 1:11:44  lr: 0.000001  loss: 7.5991 (7.8487)  task_loss: 4.7749 (5.0244)  bpp_loss: 28.2424 (28.2427)  patch_loss: 2.3215 (2.3584)  token_loss: 0.6780 (0.6676)  fea_loss: 1.7759 (1.9984)  time: 0.9647  data: 0.0001  max mem: 11194
+2024-07-25 08:52:00,570 [INFO ]  Epoch: [0]  [ 120/4448]  eta: 1:11:29  lr: 0.000001  loss: 7.5780 (7.8257)  task_loss: 4.7538 (5.0014)  bpp_loss: 28.2425 (28.2427)  patch_loss: 2.3218 (2.3558)  token_loss: 0.6774 (0.6677)  fea_loss: 1.7559 (1.9780)  time: 0.9728  data: 0.0001  max mem: 11194
+2024-07-25 08:52:10,528 [INFO ]  Epoch: [0]  [ 130/4448]  eta: 1:11:20  lr: 0.000001  loss: 7.5670 (7.8045)  task_loss: 4.7427 (4.9802)  bpp_loss: 28.2426 (28.2427)  patch_loss: 2.3239 (2.3528)  token_loss: 0.6669 (0.6682)  fea_loss: 1.7497 (1.9592)  time: 0.9861  data: 0.0001  max mem: 11194
+2024-07-25 08:52:20,550 [INFO ]  Epoch: [0]  [ 140/4448]  eta: 1:11:14  lr: 0.000001  loss: 7.5367 (7.7836)  task_loss: 4.7125 (4.9594)  bpp_loss: 28.2423 (28.2427)  patch_loss: 2.3120 (2.3499)  token_loss: 0.6791 (0.6681)  fea_loss: 1.7283 (1.9414)  time: 0.9989  data: 0.0001  max mem: 11194
+2024-07-25 08:52:30,565 [INFO ]  Epoch: [0]  [ 150/4448]  eta: 1:11:06  lr: 0.000001  loss: 7.4978 (7.7638)  task_loss: 4.6735 (4.9396)  bpp_loss: 28.2423 (28.2426)  patch_loss: 2.3131 (2.3477)  token_loss: 0.6641 (0.6675)  fea_loss: 1.6953 (1.9244)  time: 1.0018  data: 0.0001  max mem: 11194
+2024-07-25 08:52:40,399 [INFO ]  Epoch: [0]  [ 160/4448]  eta: 1:10:54  lr: 0.000001  loss: 7.4593 (7.7455)  task_loss: 4.6350 (4.9213)  bpp_loss: 28.2423 (28.2426)  patch_loss: 2.3100 (2.3450)  token_loss: 0.6513 (0.6671)  fea_loss: 1.6766 (1.9092)  time: 0.9924  data: 0.0001  max mem: 11194
+2024-07-25 08:52:50,157 [INFO ]  Epoch: [0]  [ 170/4448]  eta: 1:10:40  lr: 0.000001  loss: 7.4549 (7.7282)  task_loss: 4.6306 (4.9040)  bpp_loss: 28.2425 (28.2426)  patch_loss: 2.3005 (2.3425)  token_loss: 0.6528 (0.6669)  fea_loss: 1.6654 (1.8945)  time: 0.9795  data: 0.0001  max mem: 11194
+2024-07-25 08:52:59,911 [INFO ]  Epoch: [0]  [ 180/4448]  eta: 1:10:26  lr: 0.000001  loss: 7.4523 (7.7123)  task_loss: 4.6281 (4.8880)  bpp_loss: 28.2424 (28.2426)  patch_loss: 2.3051 (2.3407)  token_loss: 0.6546 (0.6663)  fea_loss: 1.6590 (1.8810)  time: 0.9755  data: 0.0001  max mem: 11194
+2024-07-25 08:53:09,637 [INFO ]  Epoch: [0]  [ 190/4448]  eta: 1:10:12  lr: 0.000001  loss: 7.4077 (7.6963)  task_loss: 4.5834 (4.8720)  bpp_loss: 28.2424 (28.2426)  patch_loss: 2.3067 (2.3386)  token_loss: 0.6544 (0.6660)  fea_loss: 1.6313 (1.8674)  time: 0.9739  data: 0.0001  max mem: 11194
+2024-07-25 08:53:19,360 [INFO ]  Epoch: [0]  [ 200/4448]  eta: 1:09:59  lr: 0.000001  loss: 7.4077 (7.6829)  task_loss: 4.5834 (4.8586)  bpp_loss: 28.2423 (28.2426)  patch_loss: 2.3058 (2.3366)  token_loss: 0.6651 (0.6670)  fea_loss: 1.6217 (1.8550)  time: 0.9723  data: 0.0001  max mem: 11194
+2024-07-25 08:53:29,073 [INFO ]  Epoch: [0]  [ 210/4448]  eta: 1:09:45  lr: 0.000001  loss: 7.4075 (7.6695)  task_loss: 4.5833 (4.8453)  bpp_loss: 28.2422 (28.2425)  patch_loss: 2.2973 (2.3348)  token_loss: 0.6806 (0.6674)  fea_loss: 1.6116 (1.8431)  time: 0.9717  data: 0.0001  max mem: 11195
+2024-07-25 08:53:38,788 [INFO ]  Epoch: [0]  [ 220/4448]  eta: 1:09:32  lr: 0.000001  loss: 7.3916 (7.6572)  task_loss: 4.5674 (4.8329)  bpp_loss: 28.2419 (28.2425)  patch_loss: 2.2990 (2.3332)  token_loss: 0.6721 (0.6680)  fea_loss: 1.5945 (1.8317)  time: 0.9713  data: 0.0001  max mem: 11195
+2024-07-25 08:53:48,511 [INFO ]  Epoch: [0]  [ 230/4448]  eta: 1:09:20  lr: 0.000001  loss: 7.3755 (7.6451)  task_loss: 4.5514 (4.8209)  bpp_loss: 28.2417 (28.2425)  patch_loss: 2.2998 (2.3316)  token_loss: 0.6759 (0.6679)  fea_loss: 1.5874 (1.8213)  time: 0.9718  data: 0.0001  max mem: 11195
+2024-07-25 08:53:58,215 [INFO ]  Epoch: [0]  [ 240/4448]  eta: 1:09:07  lr: 0.000001  loss: 7.3750 (7.6335)  task_loss: 4.5508 (4.8092)  bpp_loss: 28.2416 (28.2425)  patch_loss: 2.3008 (2.3302)  token_loss: 0.6651 (0.6679)  fea_loss: 1.5759 (1.8111)  time: 0.9713  data: 0.0001  max mem: 11195
+2024-07-25 08:54:07,932 [INFO ]  Epoch: [0]  [ 250/4448]  eta: 1:08:55  lr: 0.000001  loss: 7.3393 (7.6212)  task_loss: 4.5152 (4.7970)  bpp_loss: 28.2414 (28.2424)  patch_loss: 2.2951 (2.3289)  token_loss: 0.6563 (0.6671)  fea_loss: 1.5655 (1.8010)  time: 0.9710  data: 0.0001  max mem: 11195
+2024-07-25 08:54:17,651 [INFO ]  Epoch: [0]  [ 260/4448]  eta: 1:08:43  lr: 0.000001  loss: 7.3268 (7.6096)  task_loss: 4.5027 (4.7854)  bpp_loss: 28.2413 (28.2424)  patch_loss: 2.2931 (2.3274)  token_loss: 0.6561 (0.6665)  fea_loss: 1.5615 (1.7915)  time: 0.9717  data: 0.0001  max mem: 11195
+2024-07-25 08:54:27,368 [INFO ]  Epoch: [0]  [ 270/4448]  eta: 1:08:31  lr: 0.000001  loss: 7.3396 (7.5996)  task_loss: 4.5156 (4.7754)  bpp_loss: 28.2411 (28.2423)  patch_loss: 2.2837 (2.3257)  token_loss: 0.6719 (0.6672)  fea_loss: 1.5560 (1.7825)  time: 0.9717  data: 0.0001  max mem: 11195
+2024-07-25 08:54:37,070 [INFO ]  Epoch: [0]  [ 280/4448]  eta: 1:08:19  lr: 0.000001  loss: 7.3154 (7.5896)  task_loss: 4.4914 (4.7654)  bpp_loss: 28.2404 (28.2422)  patch_loss: 2.2837 (2.3246)  token_loss: 0.6619 (0.6668)  fea_loss: 1.5447 (1.7740)  time: 0.9709  data: 0.0001  max mem: 11195
+2024-07-25 08:54:46,786 [INFO ]  Epoch: [0]  [ 290/4448]  eta: 1:08:08  lr: 0.000001  loss: 7.3109 (7.5802)  task_loss: 4.4868 (4.7560)  bpp_loss: 28.2396 (28.2421)  patch_loss: 2.2924 (2.3236)  token_loss: 0.6524 (0.6669)  fea_loss: 1.5375 (1.7655)  time: 0.9708  data: 0.0001  max mem: 11195
+2024-07-25 08:54:56,541 [INFO ]  Epoch: [0]  [ 300/4448]  eta: 1:07:57  lr: 0.000001  loss: 7.3166 (7.5713)  task_loss: 4.4926 (4.7471)  bpp_loss: 28.2385 (28.2420)  patch_loss: 2.2950 (2.3225)  token_loss: 0.6756 (0.6672)  fea_loss: 1.5210 (1.7574)  time: 0.9735  data: 0.0001  max mem: 11195
+2024-07-25 08:55:06,276 [INFO ]  Epoch: [0]  [ 310/4448]  eta: 1:07:46  lr: 0.000001  loss: 7.3026 (7.5628)  task_loss: 4.4788 (4.7386)  bpp_loss: 28.2378 (28.2418)  patch_loss: 2.2952 (2.3217)  token_loss: 0.6699 (0.6673)  fea_loss: 1.5175 (1.7496)  time: 0.9744  data: 0.0001  max mem: 11195
+2024-07-25 08:55:15,984 [INFO ]  Epoch: [0]  [ 320/4448]  eta: 1:07:34  lr: 0.000001  loss: 7.2847 (7.5538)  task_loss: 4.4616 (4.7296)  bpp_loss: 28.2320 (28.2415)  patch_loss: 2.2890 (2.3206)  token_loss: 0.6612 (0.6672)  fea_loss: 1.5070 (1.7418)  time: 0.9720  data: 0.0001  max mem: 11195
+2024-07-25 08:55:25,715 [INFO ]  Epoch: [0]  [ 330/4448]  eta: 1:07:23  lr: 0.000001  loss: 7.2884 (7.5463)  task_loss: 4.4653 (4.7222)  bpp_loss: 28.2311 (28.2412)  patch_loss: 2.2923 (2.3199)  token_loss: 0.6727 (0.6678)  fea_loss: 1.4998 (1.7344)  time: 0.9719  data: 0.0001  max mem: 11195
+2024-07-25 08:55:35,470 [INFO ]  Epoch: [0]  [ 340/4448]  eta: 1:07:13  lr: 0.000001  loss: 7.2818 (7.5382)  task_loss: 4.4587 (4.7141)  bpp_loss: 28.2306 (28.2408)  patch_loss: 2.2966 (2.3192)  token_loss: 0.6657 (0.6677)  fea_loss: 1.4959 (1.7273)  time: 0.9742  data: 0.0001  max mem: 11195
+2024-07-25 08:55:45,224 [INFO ]  Epoch: [0]  [ 350/4448]  eta: 1:07:02  lr: 0.000001  loss: 7.2619 (7.5306)  task_loss: 4.4388 (4.7066)  bpp_loss: 28.2131 (28.2400)  patch_loss: 2.2874 (2.3183)  token_loss: 0.6561 (0.6679)  fea_loss: 1.4892 (1.7205)  time: 0.9754  data: 0.0001  max mem: 11195
+2024-07-25 08:55:54,962 [INFO ]  Epoch: [0]  [ 360/4448]  eta: 1:06:51  lr: 0.000001  loss: 7.2578 (7.5231)  task_loss: 4.4367 (4.6991)  bpp_loss: 28.2120 (28.2392)  patch_loss: 2.2885 (2.3175)  token_loss: 0.6579 (0.6678)  fea_loss: 1.4836 (1.7138)  time: 0.9745  data: 0.0001  max mem: 11195
+2024-07-25 08:56:04,713 [INFO ]  Epoch: [0]  [ 370/4448]  eta: 1:06:41  lr: 0.000001  loss: 7.2578 (7.5154)  task_loss: 4.4367 (4.6916)  bpp_loss: 28.2115 (28.2385)  patch_loss: 2.2920 (2.3168)  token_loss: 0.6579 (0.6676)  fea_loss: 1.4752 (1.7072)  time: 0.9743  data: 0.0001  max mem: 11195
+2024-07-25 08:56:14,459 [INFO ]  Epoch: [0]  [ 380/4448]  eta: 1:06:30  lr: 0.000001  loss: 7.2504 (7.5086)  task_loss: 4.4293 (4.6848)  bpp_loss: 28.2110 (28.2377)  patch_loss: 2.2893 (2.3161)  token_loss: 0.6630 (0.6677)  fea_loss: 1.4696 (1.7010)  time: 0.9748  data: 0.0001  max mem: 11195
+2024-07-25 08:56:24,203 [INFO ]  Epoch: [0]  [ 390/4448]  eta: 1:06:20  lr: 0.000001  loss: 7.2421 (7.5016)  task_loss: 4.4212 (4.6779)  bpp_loss: 28.2100 (28.2370)  patch_loss: 2.2893 (2.3154)  token_loss: 0.6630 (0.6676)  fea_loss: 1.4607 (1.6949)  time: 0.9744  data: 0.0001  max mem: 11195
+2024-07-25 08:56:33,961 [INFO ]  Epoch: [0]  [ 400/4448]  eta: 1:06:10  lr: 0.000001  loss: 7.2267 (7.4947)  task_loss: 4.4057 (4.6711)  bpp_loss: 28.2090 (28.2363)  patch_loss: 2.2825 (2.3147)  token_loss: 0.6591 (0.6674)  fea_loss: 1.4549 (1.6890)  time: 0.9750  data: 0.0001  max mem: 11195
+2024-07-25 08:56:43,733 [INFO ]  Epoch: [0]  [ 410/4448]  eta: 1:05:59  lr: 0.000001  loss: 7.2241 (7.4887)  task_loss: 4.4032 (4.6651)  bpp_loss: 28.2080 (28.2356)  patch_loss: 2.2854 (2.3141)  token_loss: 0.6674 (0.6676)  fea_loss: 1.4618 (1.6834)  time: 0.9764  data: 0.0001  max mem: 11195
+2024-07-25 08:56:53,446 [INFO ]  Epoch: [0]  [ 420/4448]  eta: 1:05:49  lr: 0.000001  loss: 7.2331 (7.4827)  task_loss: 4.4122 (4.6592)  bpp_loss: 28.2068 (28.2349)  patch_loss: 2.2857 (2.3134)  token_loss: 0.6711 (0.6679)  fea_loss: 1.4545 (1.6779)  time: 0.9742  data: 0.0001  max mem: 11195
+2024-07-25 08:57:03,201 [INFO ]  Epoch: [0]  [ 430/4448]  eta: 1:05:38  lr: 0.000001  loss: 7.2254 (7.4765)  task_loss: 4.4047 (4.6531)  bpp_loss: 28.2055 (28.2342)  patch_loss: 2.2863 (2.3127)  token_loss: 0.6718 (0.6679)  fea_loss: 1.4462 (1.6725)  time: 0.9733  data: 0.0001  max mem: 11195
+2024-07-25 08:57:12,938 [INFO ]  Epoch: [0]  [ 440/4448]  eta: 1:05:28  lr: 0.000001  loss: 7.2275 (7.4714)  task_loss: 4.4071 (4.6480)  bpp_loss: 28.2046 (28.2335)  patch_loss: 2.2863 (2.3122)  token_loss: 0.6750 (0.6683)  fea_loss: 1.4481 (1.6675)  time: 0.9745  data: 0.0001  max mem: 11195
+2024-07-25 08:57:22,692 [INFO ]  Epoch: [0]  [ 450/4448]  eta: 1:05:18  lr: 0.000001  loss: 7.2270 (7.4658)  task_loss: 4.4067 (4.6425)  bpp_loss: 28.2033 (28.2329)  patch_loss: 2.2936 (2.3118)  token_loss: 0.6782 (0.6684)  fea_loss: 1.4495 (1.6624)  time: 0.9745  data: 0.0001  max mem: 11195
+2024-07-25 08:57:32,428 [INFO ]  Epoch: [0]  [ 460/4448]  eta: 1:05:07  lr: 0.000001  loss: 7.2030 (7.4601)  task_loss: 4.3827 (4.6368)  bpp_loss: 28.2019 (28.2322)  patch_loss: 2.2834 (2.3110)  token_loss: 0.6762 (0.6684)  fea_loss: 1.4385 (1.6574)  time: 0.9744  data: 0.0001  max mem: 11195
+2024-07-25 08:57:42,172 [INFO ]  Epoch: [0]  [ 470/4448]  eta: 1:04:57  lr: 0.000001  loss: 7.2220 (7.4547)  task_loss: 4.4020 (4.6315)  bpp_loss: 28.2006 (28.2315)  patch_loss: 2.2808 (2.3105)  token_loss: 0.6535 (0.6685)  fea_loss: 1.4261 (1.6526)  time: 0.9739  data: 0.0001  max mem: 11195
+2024-07-25 08:57:51,902 [INFO ]  Epoch: [0]  [ 480/4448]  eta: 1:04:47  lr: 0.000001  loss: 7.2008 (7.4490)  task_loss: 4.3808 (4.6259)  bpp_loss: 28.1997 (28.2308)  patch_loss: 2.2765 (2.3099)  token_loss: 0.6534 (0.6682)  fea_loss: 1.4257 (1.6478)  time: 0.9736  data: 0.0001  max mem: 11195
+2024-07-25 08:58:01,653 [INFO ]  Epoch: [0]  [ 490/4448]  eta: 1:04:37  lr: 0.000001  loss: 7.2050 (7.4444)  task_loss: 4.3852 (4.6214)  bpp_loss: 28.1985 (28.2301)  patch_loss: 2.2831 (2.3095)  token_loss: 0.6743 (0.6686)  fea_loss: 1.4232 (1.6432)  time: 0.9740  data: 0.0001  max mem: 11195
+2024-07-25 08:58:11,373 [INFO ]  Epoch: [0]  [ 500/4448]  eta: 1:04:26  lr: 0.000001  loss: 7.2024 (7.4392)  task_loss: 4.3827 (4.6162)  bpp_loss: 28.1970 (28.2295)  patch_loss: 2.2877 (2.3091)  token_loss: 0.6743 (0.6684)  fea_loss: 1.4212 (1.6387)  time: 0.9735  data: 0.0001  max mem: 11195
+2024-07-25 08:58:21,089 [INFO ]  Epoch: [0]  [ 510/4448]  eta: 1:04:16  lr: 0.000001  loss: 7.1806 (7.4341)  task_loss: 4.3611 (4.6112)  bpp_loss: 28.1960 (28.2288)  patch_loss: 2.2865 (2.3086)  token_loss: 0.6606 (0.6683)  fea_loss: 1.4148 (1.6343)  time: 0.9717  data: 0.0001  max mem: 11195
+2024-07-25 08:58:30,784 [INFO ]  Epoch: [0]  [ 520/4448]  eta: 1:04:05  lr: 0.000001  loss: 7.1758 (7.4292)  task_loss: 4.3564 (4.6064)  bpp_loss: 28.1949 (28.2282)  patch_loss: 2.2844 (2.3081)  token_loss: 0.6636 (0.6683)  fea_loss: 1.4082 (1.6299)  time: 0.9705  data: 0.0001  max mem: 11195
+2024-07-25 08:58:40,522 [INFO ]  Epoch: [0]  [ 530/4448]  eta: 1:03:55  lr: 0.000001  loss: 7.1598 (7.4244)  task_loss: 4.3403 (4.6017)  bpp_loss: 28.1940 (28.2275)  patch_loss: 2.2851 (2.3077)  token_loss: 0.6642 (0.6683)  fea_loss: 1.4067 (1.6258)  time: 0.9715  data: 0.0001  max mem: 11195
+2024-07-25 08:58:50,244 [INFO ]  Epoch: [0]  [ 540/4448]  eta: 1:03:45  lr: 0.000001  loss: 7.1524 (7.4195)  task_loss: 4.3331 (4.5969)  bpp_loss: 28.1934 (28.2269)  patch_loss: 2.2790 (2.3072)  token_loss: 0.6654 (0.6681)  fea_loss: 1.4014 (1.6216)  time: 0.9729  data: 0.0001  max mem: 11195
+2024-07-25 08:59:00,201 [INFO ]  Epoch: [0]  [ 550/4448]  eta: 1:03:36  lr: 0.000001  loss: 7.1595 (7.4148)  task_loss: 4.3401 (4.5922)  bpp_loss: 28.1927 (28.2262)  patch_loss: 2.2758 (2.3066)  token_loss: 0.6682 (0.6682)  fea_loss: 1.3959 (1.6174)  time: 0.9839  data: 0.0001  max mem: 11195

1_feature_extractor/log/DINOv2_training/log/20240730_084148.log

@@ -0,0 +1,301 @@
+2024-07-30 08:41:48,409 [INFO ]  Logging file is /data0/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/log//20240730_084148.log
+2024-07-30 08:41:48,410 [INFO ]  job dir: /data0/qiyp/Proteus-pytorch/pretrain
+2024-07-30 08:41:48,410 [INFO ]  Namespace(batch_size=48,
+epochs=200,
+bce_loss=False,
+unscale_lr=False,
+model='models_proteus_dinov2',
+target_model='vit_base',
+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.0002,
+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,
+global_crops_size=224,
+patch_size=14,
+mask_ratio=[0.5],
+mask_probability=0.5,
+mask_first_n=True,
+clone_batch=1,
+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='vit_large',
+teacher_path='',
+distillation_type='none',
+distillation_alpha=0.5,
+distillation_tau=1.0,
+lambda_token=1.0,
+lambda_fea=1.1,
+lambda_patch=1.0,
+cosub=False,
+finetune='',
+attn_only=False,
+weight_inherit='',
+data_path='/data1/datasets/imagenet_fold',
+data_set='IMNET',
+inat_category='name',
+output_dir='log/DINOv2_training',
+log_dir='/data0/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/log/',
+device='cuda',
+seed=0,
+resume='/data0/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/checkpoint0080.pth',
+start_epoch=0,
+eval=False,
+eval_crop_ratio=0.875,
+dist_eval=False,
+num_workers=10,
+pin_mem=True,
+distributed=True,
+world_size=6,
+dist_url='env://',
+rank=0,
+gpu=0,
+dist_backend='nccl')
+
+2024-07-30 08:41:52,685 [INFO ]  Dataset ImageFolder
+    Number of datapoints: 1281167
+    Root location: /data1/datasets/imagenet_fold/train
+    StandardTransform
+Transform: Compose(
+               RandomResizedCropAndInterpolation(size=(224, 224), scale=(0.08, 1.0), ratio=(0.75, 1.3333), interpolation=bicubic)
+               RandomHorizontalFlip(p=0.5)
+               RandAugment(n=2, ops=
+           	AugmentOp(name=AutoContrast, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=Equalize, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=Invert, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=Rotate, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=PosterizeIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=SolarizeIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=SolarizeAdd, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=ColorIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=ContrastIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=BrightnessIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=SharpnessIncreasing, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=ShearX, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=ShearY, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=TranslateXRel, p=0.5, m=9, mstd=0.5)
+           	AugmentOp(name=TranslateYRel, p=0.5, m=9, mstd=0.5))
+               ToTensor()
+               Normalize(mean=tensor([0.4850, 0.4560, 0.4060]), std=tensor([0.2290, 0.2240, 0.2250]))
+               RandomErasing(p=0.25, mode=pixel, count=(1, 1))
+           )
+2024-07-30 08:41:52,760 [INFO ]  Sampler_train = <samplers.RASampler object at 0x7efd92aaf3a0>
+2024-07-30 08:41:52,792 [INFO ]  Created a temporary directory at /tmp/tmp1r345n4h
+2024-07-30 08:41:52,793 [INFO ]  Writing /tmp/tmp1r345n4h/_remote_module_non_scriptable.py
+2024-07-30 08:42:03,348 [INFO ]  using MLP layer as FFN
+2024-07-30 08:42:05,496 [INFO ]  using MLP layer as FFN
+2024-07-30 08:42:13,281 [INFO ]  Model = MetaArch(
+  (student): ModuleDict(
+    (backbone): DinoVisionTransformer(
+      (patch_embed): PatchEmbed(
+        (proj): Conv2d(3, 768, kernel_size=(14, 14), stride=(14, 14))
+        (norm): Identity()
+      )
+      (blocks): ModuleList(
+        (0-11): 12 x Block(
+          (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
+          (attn): MemEffAttention(
+            (qkv): Linear(in_features=768, out_features=2304, bias=True)
+            (attn_drop): Dropout(p=0.0, inplace=False)
+            (proj): Linear(in_features=768, out_features=768, bias=True)
+            (proj_drop): Dropout(p=0.0, inplace=False)
+          )
+          (ls1): LayerScale()
+          (drop_path1): Identity()
+          (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
+          (mlp): Mlp(
+            (fc1): Linear(in_features=768, out_features=3072, bias=True)
+            (act): GELU(approximate='none')
+            (fc2): Linear(in_features=3072, out_features=768, bias=True)
+            (drop): Dropout(p=0.0, inplace=False)
+          )
+          (ls2): LayerScale()
+          (drop_path2): Identity()
+        )
+      )
+      (norm): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
+      (head): Identity()
+    )
+  )
+  (teacher): ModuleDict(
+    (backbone): DinoVisionTransformer(
+      (patch_embed): PatchEmbed(
+        (proj): Conv2d(3, 1024, kernel_size=(14, 14), stride=(14, 14))
+        (norm): Identity()
+      )
+      (blocks): ModuleList(
+        (0-23): 24 x NestedTensorBlock(
+          (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
+          (attn): MemEffAttention(
+            (qkv): Linear(in_features=1024, out_features=3072, bias=True)
+            (attn_drop): Dropout(p=0.0, inplace=False)
+            (proj): Linear(in_features=1024, out_features=1024, bias=True)
+            (proj_drop): Dropout(p=0.0, inplace=False)
+          )
+          (ls1): LayerScale()
+          (drop_path1): Identity()
+          (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
+          (mlp): Mlp(
+            (fc1): Linear(in_features=1024, out_features=4096, bias=True)
+            (act): GELU(approximate='none')
+            (fc2): Linear(in_features=4096, out_features=1024, bias=True)
+            (drop): Dropout(p=0.0, inplace=False)
+          )
+          (ls2): LayerScale()
+          (drop_path2): Identity()
+        )
+      )
+      (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
+      (head): Identity()
+    )
+  )
+  (ibot_head): Sequential(
+    (0): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (1): Linear(in_features=768, out_features=1024, bias=True)
+  )
+  (token_head): Sequential(
+    (0): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (1): Linear(in_features=768, out_features=1024, bias=True)
+  )
+  (fea_head): Sequential(
+    (0): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (1): Linear(in_features=768, out_features=1024, bias=True)
+  )
+  (soft_criterion): MSELoss()
+  (info_bottleneck): IF_Module(
+    (encoder_blocks): ModuleList(
+      (0-3): 4 x Block(
+        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+        (attn): Attention(
+          (qkv): Linear(in_features=768, out_features=2304, bias=True)
+          (q_norm): Identity()
+          (k_norm): Identity()
+          (attn_drop): Dropout(p=0.0, inplace=False)
+          (proj): Linear(in_features=768, out_features=768, bias=True)
+          (proj_drop): Dropout(p=0.0, inplace=False)
+        )
+        (ls1): Identity()
+        (drop_path1): Identity()
+        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+        (mlp): Mlp(
+          (fc1): Linear(in_features=768, out_features=3072, bias=True)
+          (act): GELU(approximate='none')
+          (drop1): Dropout(p=0.0, inplace=False)
+          (norm): Identity()
+          (fc2): Linear(in_features=3072, out_features=768, bias=True)
+          (drop2): Dropout(p=0.0, inplace=False)
+        )
+        (ls2): Identity()
+        (drop_path2): Identity()
+      )
+    )
+    (encoder_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (decoder_blocks): ModuleList(
+      (0-3): 4 x Block(
+        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+        (attn): Attention(
+          (qkv): Linear(in_features=768, out_features=2304, bias=True)
+          (q_norm): Identity()
+          (k_norm): Identity()
+          (attn_drop): Dropout(p=0.0, inplace=False)
+          (proj): Linear(in_features=768, out_features=768, bias=True)
+          (proj_drop): Dropout(p=0.0, inplace=False)
+        )
+        (ls1): Identity()
+        (drop_path1): Identity()
+        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+        (mlp): Mlp(
+          (fc1): Linear(in_features=768, out_features=3072, bias=True)
+          (act): GELU(approximate='none')
+          (drop1): Dropout(p=0.0, inplace=False)
+          (norm): Identity()
+          (fc2): Linear(in_features=3072, out_features=768, bias=True)
+          (drop2): Dropout(p=0.0, inplace=False)
+        )
+        (ls2): Identity()
+        (drop_path2): Identity()
+      )
+    )
+    (decoder_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
+    (entropy_bottleneck): EntropyBottleneck(
+      (likelihood_lower_bound): LowerBound()
+    )
+  )
+)
+2024-07-30 08:42:16,057 [INFO ]  number of params: 144845568
+2024-07-30 08:42:16,057 [INFO ]  base lr = 0.0002
+2024-07-30 08:42:16,058 [INFO ]  actural lr = 0.00011250000000000001
+2024-07-30 08:42:35,018 [INFO ]  Loaded state_dict_ema
+2024-07-30 08:42:35,076 [INFO ]  Resuming from /data0/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/checkpoint0080.pth
+2024-07-30 08:42:35,076 [INFO ]  Start training for 200 epochs
+2024-07-30 08:42:47,484 [INFO ]  Epoch: [81]  [   0/4448]  eta: 15:19:18  lr: 0.000076  loss: 2.8205 (2.8205)  task_loss: 2.6200 (2.6200)  bpp_loss: 0.2883 (0.2883)  patch_loss: 0.9649 (0.9649)  token_loss: 0.8029 (0.8029)  fea_loss: 0.8522 (0.8522)  time: 12.4007  data: 3.4410  max mem: 20527
+2024-07-30 08:43:06,985 [INFO ]  Epoch: [81]  [  20/4448]  eta: 1:52:06  lr: 0.000076  loss: 2.8058 (2.8015)  task_loss: 2.5992 (2.5967)  bpp_loss: 0.2986 (0.2988)  patch_loss: 0.9937 (0.9960)  token_loss: 0.7355 (0.7480)  fea_loss: 0.8571 (0.8527)  time: 0.9749  data: 0.0001  max mem: 21093
+2024-07-30 08:43:26,649 [INFO ]  Epoch: [81]  [  40/4448]  eta: 1:32:23  lr: 0.000076  loss: 2.7488 (2.7856)  task_loss: 2.5421 (2.5804)  bpp_loss: 0.3020 (0.3006)  patch_loss: 0.9832 (0.9945)  token_loss: 0.7257 (0.7364)  fea_loss: 0.8415 (0.8495)  time: 0.9831  data: 0.0001  max mem: 21095
+2024-07-30 08:43:46,729 [INFO ]  Epoch: [81]  [  60/4448]  eta: 1:25:53  lr: 0.000076  loss: 2.7729 (2.7857)  task_loss: 2.5642 (2.5803)  bpp_loss: 0.3040 (0.3017)  patch_loss: 1.0076 (0.9998)  token_loss: 0.7200 (0.7333)  fea_loss: 0.8399 (0.8472)  time: 1.0040  data: 0.0001  max mem: 21095
+2024-07-30 08:44:07,288 [INFO ]  Epoch: [81]  [  80/4448]  eta: 1:22:51  lr: 0.000076  loss: 2.8271 (2.7941)  task_loss: 2.6186 (2.5881)  bpp_loss: 0.3040 (0.3025)  patch_loss: 1.0128 (1.0035)  token_loss: 0.7397 (0.7346)  fea_loss: 0.8573 (0.8500)  time: 1.0279  data: 0.0001  max mem: 21095
+2024-07-30 08:44:27,381 [INFO ]  Epoch: [81]  [ 100/4448]  eta: 1:20:34  lr: 0.000076  loss: 2.8005 (2.8008)  task_loss: 2.5913 (2.5942)  bpp_loss: 0.3074 (0.3034)  patch_loss: 1.0188 (1.0078)  token_loss: 0.7295 (0.7347)  fea_loss: 0.8552 (0.8517)  time: 1.0046  data: 0.0001  max mem: 21095
+2024-07-30 08:44:47,266 [INFO ]  Epoch: [81]  [ 120/4448]  eta: 1:18:47  lr: 0.000076  loss: 2.7752 (2.7995)  task_loss: 2.5660 (2.5928)  bpp_loss: 0.3070 (0.3040)  patch_loss: 0.9996 (1.0060)  token_loss: 0.7348 (0.7359)  fea_loss: 0.8485 (0.8509)  time: 0.9942  data: 0.0001  max mem: 21095
+2024-07-30 08:45:07,249 [INFO ]  Epoch: [81]  [ 140/4448]  eta: 1:17:28  lr: 0.000076  loss: 2.8033 (2.8006)  task_loss: 2.5930 (2.5937)  bpp_loss: 0.3064 (0.3044)  patch_loss: 1.0073 (1.0067)  token_loss: 0.7371 (0.7357)  fea_loss: 0.8503 (0.8513)  time: 0.9991  data: 0.0002  max mem: 21096
+2024-07-30 08:45:27,250 [INFO ]  Epoch: [81]  [ 160/4448]  eta: 1:16:25  lr: 0.000076  loss: 2.7980 (2.8014)  task_loss: 2.5895 (2.5944)  bpp_loss: 0.3069 (0.3047)  patch_loss: 1.0087 (1.0062)  token_loss: 0.7382 (0.7361)  fea_loss: 0.8631 (0.8520)  time: 1.0000  data: 0.0001  max mem: 21096
+2024-07-30 08:45:47,340 [INFO ]  Epoch: [81]  [ 180/4448]  eta: 1:15:33  lr: 0.000076  loss: 2.8292 (2.8018)  task_loss: 2.6184 (2.5945)  bpp_loss: 0.3091 (0.3052)  patch_loss: 1.0158 (1.0076)  token_loss: 0.7273 (0.7351)  fea_loss: 0.8446 (0.8518)  time: 1.0044  data: 0.0001  max mem: 21096
+2024-07-30 08:46:07,943 [INFO ]  Epoch: [81]  [ 200/4448]  eta: 1:14:58  lr: 0.000076  loss: 2.7974 (2.7999)  task_loss: 2.5896 (2.5926)  bpp_loss: 0.3082 (0.3055)  patch_loss: 0.9949 (1.0069)  token_loss: 0.7182 (0.7340)  fea_loss: 0.8525 (0.8516)  time: 1.0301  data: 0.0001  max mem: 21096
+2024-07-30 08:46:28,477 [INFO ]  Epoch: [81]  [ 220/4448]  eta: 1:14:24  lr: 0.000076  loss: 2.7864 (2.8008)  task_loss: 2.5794 (2.5933)  bpp_loss: 0.3078 (0.3057)  patch_loss: 1.0100 (1.0072)  token_loss: 0.7306 (0.7342)  fea_loss: 0.8441 (0.8518)  time: 1.0266  data: 0.0001  max mem: 21096
+2024-07-30 08:46:48,558 [INFO ]  Epoch: [81]  [ 240/4448]  eta: 1:13:45  lr: 0.000076  loss: 2.8172 (2.8027)  task_loss: 2.6068 (2.5950)  bpp_loss: 0.3095 (0.3060)  patch_loss: 1.0120 (1.0081)  token_loss: 0.7279 (0.7344)  fea_loss: 0.8583 (0.8525)  time: 1.0040  data: 0.0001  max mem: 21096
+2024-07-30 08:47:08,665 [INFO ]  Epoch: [81]  [ 260/4448]  eta: 1:13:09  lr: 0.000076  loss: 2.7650 (2.8013)  task_loss: 2.5573 (2.5936)  bpp_loss: 0.3086 (0.3063)  patch_loss: 0.9966 (1.0075)  token_loss: 0.7286 (0.7338)  fea_loss: 0.8450 (0.8523)  time: 1.0053  data: 0.0001  max mem: 21096
+2024-07-30 08:47:28,757 [INFO ]  Epoch: [81]  [ 280/4448]  eta: 1:12:35  lr: 0.000076  loss: 2.7513 (2.7988)  task_loss: 2.5452 (2.5910)  bpp_loss: 0.3100 (0.3066)  patch_loss: 0.9790 (1.0058)  token_loss: 0.7214 (0.7333)  fea_loss: 0.8448 (0.8518)  time: 1.0045  data: 0.0001  max mem: 21096
+2024-07-30 08:47:48,814 [INFO ]  Epoch: [81]  [ 300/4448]  eta: 1:12:03  lr: 0.000076  loss: 2.7910 (2.7984)  task_loss: 2.5833 (2.5906)  bpp_loss: 0.3097 (0.3068)  patch_loss: 0.9979 (1.0056)  token_loss: 0.7220 (0.7334)  fea_loss: 0.8453 (0.8515)  time: 1.0028  data: 0.0001  max mem: 21096
+2024-07-30 08:48:08,893 [INFO ]  Epoch: [81]  [ 320/4448]  eta: 1:11:32  lr: 0.000076  loss: 2.7957 (2.7991)  task_loss: 2.5868 (2.5911)  bpp_loss: 0.3101 (0.3071)  patch_loss: 1.0096 (1.0065)  token_loss: 0.7125 (0.7326)  fea_loss: 0.8559 (0.8520)  time: 1.0039  data: 0.0001  max mem: 21096
+2024-07-30 08:48:29,072 [INFO ]  Epoch: [81]  [ 340/4448]  eta: 1:11:04  lr: 0.000076  loss: 2.7660 (2.7990)  task_loss: 2.5585 (2.5909)  bpp_loss: 0.3108 (0.3072)  patch_loss: 0.9918 (1.0063)  token_loss: 0.7228 (0.7327)  fea_loss: 0.8468 (0.8519)  time: 1.0089  data: 0.0001  max mem: 21096
+2024-07-30 08:48:49,101 [INFO ]  Epoch: [81]  [ 360/4448]  eta: 1:10:35  lr: 0.000076  loss: 2.8059 (2.7999)  task_loss: 2.5933 (2.5917)  bpp_loss: 0.3136 (0.3076)  patch_loss: 1.0057 (1.0066)  token_loss: 0.7353 (0.7330)  fea_loss: 0.8496 (0.8521)  time: 1.0014  data: 0.0001  max mem: 21096
+2024-07-30 08:49:09,150 [INFO ]  Epoch: [81]  [ 380/4448]  eta: 1:10:07  lr: 0.000076  loss: 2.7807 (2.7979)  task_loss: 2.5695 (2.5897)  bpp_loss: 0.3108 (0.3078)  patch_loss: 0.9834 (1.0055)  token_loss: 0.7254 (0.7327)  fea_loss: 0.8475 (0.8515)  time: 1.0024  data: 0.0001  max mem: 21096
+2024-07-30 08:49:29,184 [INFO ]  Epoch: [81]  [ 400/4448]  eta: 1:09:40  lr: 0.000076  loss: 2.7576 (2.7964)  task_loss: 2.5457 (2.5881)  bpp_loss: 0.3135 (0.3081)  patch_loss: 0.9908 (1.0050)  token_loss: 0.7106 (0.7318)  fea_loss: 0.8434 (0.8513)  time: 1.0016  data: 0.0001  max mem: 21096
+2024-07-30 08:49:49,309 [INFO ]  Epoch: [81]  [ 420/4448]  eta: 1:09:14  lr: 0.000076  loss: 2.7548 (2.7950)  task_loss: 2.5465 (2.5866)  bpp_loss: 0.3124 (0.3083)  patch_loss: 0.9950 (1.0045)  token_loss: 0.7273 (0.7312)  fea_loss: 0.8426 (0.8510)  time: 1.0062  data: 0.0002  max mem: 21096
+2024-07-30 08:50:09,308 [INFO ]  Epoch: [81]  [ 440/4448]  eta: 1:08:47  lr: 0.000076  loss: 2.7648 (2.7940)  task_loss: 2.5540 (2.5855)  bpp_loss: 0.3125 (0.3085)  patch_loss: 0.9801 (1.0036)  token_loss: 0.7217 (0.7311)  fea_loss: 0.8493 (0.8508)  time: 0.9999  data: 0.0002  max mem: 21096
+2024-07-30 08:50:29,354 [INFO ]  Epoch: [81]  [ 460/4448]  eta: 1:08:22  lr: 0.000076  loss: 2.7843 (2.7932)  task_loss: 2.5758 (2.5847)  bpp_loss: 0.3126 (0.3087)  patch_loss: 0.9979 (1.0033)  token_loss: 0.7206 (0.7308)  fea_loss: 0.8484 (0.8506)  time: 1.0022  data: 0.0002  max mem: 21096
+2024-07-30 08:50:49,382 [INFO ]  Epoch: [81]  [ 480/4448]  eta: 1:07:57  lr: 0.000076  loss: 2.7398 (2.7920)  task_loss: 2.5274 (2.5834)  bpp_loss: 0.3156 (0.3090)  patch_loss: 0.9877 (1.0032)  token_loss: 0.7010 (0.7300)  fea_loss: 0.8394 (0.8502)  time: 1.0013  data: 0.0002  max mem: 21096
+2024-07-30 08:51:09,296 [INFO ]  Epoch: [81]  [ 500/4448]  eta: 1:07:31  lr: 0.000076  loss: 2.7482 (2.7908)  task_loss: 2.5427 (2.5822)  bpp_loss: 0.3129 (0.3091)  patch_loss: 0.9996 (1.0030)  token_loss: 0.6940 (0.7294)  fea_loss: 0.8401 (0.8499)  time: 0.9956  data: 0.0002  max mem: 21096
+2024-07-30 08:51:29,262 [INFO ]  Epoch: [81]  [ 520/4448]  eta: 1:07:07  lr: 0.000076  loss: 2.7785 (2.7908)  task_loss: 2.5668 (2.5821)  bpp_loss: 0.3141 (0.3093)  patch_loss: 1.0087 (1.0028)  token_loss: 0.7261 (0.7294)  fea_loss: 0.8467 (0.8498)  time: 0.9982  data: 0.0002  max mem: 21096
+2024-07-30 08:51:49,227 [INFO ]  Epoch: [81]  [ 540/4448]  eta: 1:06:42  lr: 0.000076  loss: 2.8121 (2.7921)  task_loss: 2.5996 (2.5833)  bpp_loss: 0.3148 (0.3096)  patch_loss: 0.9975 (1.0033)  token_loss: 0.7236 (0.7298)  fea_loss: 0.8581 (0.8502)  time: 0.9982  data: 0.0002  max mem: 21096
+2024-07-30 08:52:09,093 [INFO ]  Epoch: [81]  [ 560/4448]  eta: 1:06:17  lr: 0.000076  loss: 2.7914 (2.7926)  task_loss: 2.5781 (2.5837)  bpp_loss: 0.3144 (0.3097)  patch_loss: 1.0098 (1.0037)  token_loss: 0.7028 (0.7295)  fea_loss: 0.8557 (0.8505)  time: 0.9932  data: 0.0002  max mem: 21096
+2024-07-30 08:52:29,024 [INFO ]  Epoch: [81]  [ 580/4448]  eta: 1:05:53  lr: 0.000076  loss: 2.8253 (2.7942)  task_loss: 2.6127 (2.5851)  bpp_loss: 0.3147 (0.3099)  patch_loss: 1.0152 (1.0044)  token_loss: 0.7303 (0.7297)  fea_loss: 0.8694 (0.8511)  time: 0.9965  data: 0.0002  max mem: 21096
+2024-07-30 08:52:48,932 [INFO ]  Epoch: [81]  [ 600/4448]  eta: 1:05:30  lr: 0.000076  loss: 2.8256 (2.7949)  task_loss: 2.6144 (2.5858)  bpp_loss: 0.3141 (0.3101)  patch_loss: 0.9969 (1.0043)  token_loss: 0.7540 (0.7303)  fea_loss: 0.8563 (0.8512)  time: 0.9953  data: 0.0002  max mem: 21096
+2024-07-30 08:53:08,892 [INFO ]  Epoch: [81]  [ 620/4448]  eta: 1:05:06  lr: 0.000076  loss: 2.7807 (2.7948)  task_loss: 2.5716 (2.5856)  bpp_loss: 0.3145 (0.3102)  patch_loss: 1.0056 (1.0042)  token_loss: 0.7300 (0.7303)  fea_loss: 0.8505 (0.8510)  time: 0.9980  data: 0.0002  max mem: 21096
+2024-07-30 08:53:28,936 [INFO ]  Epoch: [81]  [ 640/4448]  eta: 1:04:44  lr: 0.000076  loss: 2.8072 (2.7954)  task_loss: 2.5994 (2.5861)  bpp_loss: 0.3138 (0.3103)  patch_loss: 1.0003 (1.0044)  token_loss: 0.7397 (0.7304)  fea_loss: 0.8647 (0.8514)  time: 1.0021  data: 0.0002  max mem: 21096
+2024-07-30 08:53:48,900 [INFO ]  Epoch: [81]  [ 660/4448]  eta: 1:04:21  lr: 0.000076  loss: 2.7934 (2.7958)  task_loss: 2.5801 (2.5865)  bpp_loss: 0.3144 (0.3104)  patch_loss: 0.9936 (1.0043)  token_loss: 0.7337 (0.7307)  fea_loss: 0.8504 (0.8515)  time: 0.9982  data: 0.0002  max mem: 21096
+2024-07-30 08:54:08,946 [INFO ]  Epoch: [81]  [ 680/4448]  eta: 1:03:58  lr: 0.000076  loss: 2.7859 (2.7956)  task_loss: 2.5747 (2.5862)  bpp_loss: 0.3140 (0.3106)  patch_loss: 0.9925 (1.0043)  token_loss: 0.7216 (0.7306)  fea_loss: 0.8422 (0.8514)  time: 1.0022  data: 0.0002  max mem: 21096
+2024-07-30 08:54:28,905 [INFO ]  Epoch: [81]  [ 700/4448]  eta: 1:03:36  lr: 0.000076  loss: 2.7273 (2.7944)  task_loss: 2.5182 (2.5851)  bpp_loss: 0.3142 (0.3107)  patch_loss: 0.9854 (1.0038)  token_loss: 0.7153 (0.7301)  fea_loss: 0.8333 (0.8511)  time: 0.9979  data: 0.0001  max mem: 21096

1_feature_extractor/losses_hint.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+"""
+Implements the knowledge distillation loss
+"""
+import torch
+from torch.nn import functional as F
+
+
+class DistillationLoss(torch.nn.Module):
+    """
+    This module wraps a standard criterion and adds an extra knowledge distillation loss by
+    taking a teacher model prediction and using it as additional supervision.
+    """
+    def __init__(self, base_criterion: torch.nn.Module, teacher_model: torch.nn.Module,
+                 distillation_type: str, lambda_token: float, lambda_fea: float):
+        super().__init__()
+        self.base_criterion = base_criterion
+        self.teacher_model = teacher_model
+        assert distillation_type in ['none', 'soft', 'hard']
+        self.distillation_type = distillation_type
+        self.lambda_token = lambda_token
+        self.lambda_fea = lambda_fea
+        self.soft_criterion = torch.nn.MSELoss()
+
+    def forward(self, inputs, outputs, labels):
+        """
+        Args:
+            inputs: The original inputs that are feed to the teacher model
+            outputs: the outputs of the model to be trained. It is expected to be
+                either a Tensor, or a Tuple[Tensor, Tensor], with the original output
+                in the first position and the distillation predictions as the second output
+            labels: the labels for the base criterion
+        """
+        outputs_token, outputs_fea = outputs
+
+        # don't backprop throught the teacher
+        with torch.no_grad():
+            teacher_outputs = self.teacher_model.backbone.get_intermediate_layers(inputs, n=4, return_class_token=True)
+            teacher_outputs_token = teacher_outputs[3][1]
+            teacher_outputs_fea = torch.cat((teacher_outputs_token.unsqueeze(1),teacher_outputs[3][0]),dim=1)
+
+        distillation_loss_token = self.soft_criterion(outputs_token, teacher_outputs_token)
+        distillation_loss_fea = self.soft_criterion(outputs_fea, teacher_outputs_fea)
+
+        token_loss = self.lambda_token * distillation_loss_token
+        fea_loss = self.lambda_fea * distillation_loss_fea
+        
+        return token_loss, fea_loss

1_feature_extractor/main.py

@@ -0,0 +1,520 @@
+# 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 timm.models import create_model
+from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
+from timm.scheduler import create_scheduler
+from timm.optim import create_optimizer
+from timm.utils import NativeScaler, get_state_dict, ModelEma
+from augmentations import collate_data_and_cast_aug
+from datasets import build_dataset
+
+from losses_hint import DistillationLoss
+from samplers import RASampler
+from functools import partial
+
+import importlib
+import utils
+import random
+import math
+from multiprocessing import Value
+from abc import ABC
+
+import sys
+from typing import Iterable, Optional
+from timm.data import Mixup
+from timm.utils import accuracy, ModelEma
+import utils
+
+
+class MaskingGenerator(ABC):
+    def __init__(self, input_size):
+        if not isinstance(input_size, tuple):
+            input_size = (input_size,) * 2
+        self.height, self.width = input_size
+        self.num_patches = self.height * self.width
+
+    def __repr__(self):
+        raise NotImplementedError
+
+    def get_shape(self):
+        return self.height, self.width
+
+    def _mask(self, mask, max_mask_patches):
+        raise NotImplementedError
+
+    def get_none_mask(self):
+        return np.zeros(shape=self.get_shape(), dtype=bool)
+    
+    
+    
+class RandomMaskingGenerator(MaskingGenerator):
+    def __init__(
+        self,
+        input_size,
+    ):
+        """
+        Args:
+            input_size: the size of the token map, e.g., 14x14
+        """
+        super().__init__(input_size)
+
+    def __repr__(self):
+        repr_str = f"Random Generator({self.height}, {self.width})"
+        return repr_str
+
+    def _mask(self, mask, max_mask_patches):
+        return super()._mask(mask, max_mask_patches)
+
+    def __call__(self, num_masking_patches=0):
+        if num_masking_patches <= 0:
+            return np.zeros(shape=self.get_shape(), dtype=bool)
+
+        mask = np.hstack([np.ones(num_masking_patches, dtype=bool),
+                          np.zeros(self.num_patches - num_masking_patches, dtype=bool)])
+        np.random.shuffle(mask)
+        mask = mask.reshape(self.get_shape())
+        return mask
+
+
+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)
+    parser.add_argument('--target_model', default='deit_base_patch16_224', type=str)
+    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
+    
+    # add dataset parameters
+    parser.add_argument('--global_crops_size', '--img_size', default=224, type=int,
+                        help="this should be equal to image size")
+    parser.add_argument('--patch_size', default=16, type=int,
+                        help="patch size for vit patch embedding")
+    
+    # add masking parameter
+    parser.add_argument('--mask_ratio', default=(0.1, 0.5), type=float, nargs='+',
+                        help="mask ratio can be either a value or a range")
+    parser.add_argument('--mask_probability', default=0., type=float,
+                        help="how many samples with be applied with masking")
+    parser.add_argument('--mask_first_n', action='store_true',
+                        help="mask the first n sample to avoid shuffling. Needed for MAE-style encoder")
+    parser.add_argument('--clone_batch', default=1, type=int,
+                        help="how many times to clone the batch for masking (default: 1, not cloning)")
+    
+    # * 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='base', type=str)
+    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="")
+    parser.add_argument('--lambda_token', type=float, default=1.0)
+    parser.add_argument('--lambda_fea', type=float, default=1.0)
+    parser.add_argument('--lambda_patch', type=float, default=1.0)
+    
+    # * Cosub params
+    parser.add_argument('--cosub', action='store_true') 
+    
+    # * Finetuning params
+    parser.add_argument('--finetune', default='', help='finetune from checkpoint')
+    parser.add_argument('--attn-only', action='store_true')
+    parser.add_argument('--weight_inherit', default='') 
+    
+    # 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', 'IMNET_ibot', 'IMNET_ibot_aug', 'IMNET_ibot_fast_aug', 'INAT', 'INAT19', 'IMNET_L', 'IMNET_L_ibot'],
+                        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='cuda',
+                        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('--distributed', action='store_true', default=False, help='Enabling distributed training')
+    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')
+    return parser
+
+
+def main(args):
+    utils.init_distributed_mode(args)
+
+    print(args)
+
+    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
+
+    dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
+
+    if args.distributed:
+        num_tasks = utils.get_world_size()
+        global_rank = utils.get_rank()
+        if args.repeated_aug:
+            sampler_train = RASampler(
+                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+            )
+        else:
+            sampler_train = torch.utils.data.DistributedSampler(
+                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+            )
+    else:
+        sampler_train = torch.utils.data.RandomSampler(dataset_train)
+        
+    n_tokens = (args.global_crops_size // args.patch_size) ** 2
+    mask_generator = RandomMaskingGenerator(
+        input_size=args.global_crops_size // args.patch_size,
+    )
+
+    collate_fn = partial(
+        collate_data_and_cast_aug,
+        mask_ratio=args.mask_ratio,
+        mask_probability=args.mask_probability,
+        dtype=torch.half,   # half precision
+        n_tokens=n_tokens,
+        mask_first_n=args.mask_first_n,
+        mask_generator=mask_generator,
+        clone_batch=args.clone_batch,
+    )
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train, sampler=sampler_train,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+        collate_fn=collate_fn,
+    )
+
+    mixup_fn = None
+
+    print(f"Creating model: {args.model}")
+    meta_arch_module = importlib.import_module(args.model)
+    MetaArch = meta_arch_module.MetaArch
+
+    model = MetaArch(args)
+
+    if args.finetune:
+        checkpoint = torch.load(args.finetune, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            pretrained_dict = checkpoint['state_dict']
+        elif 'model' in checkpoint:
+            pretrained_dict = checkpoint['model']
+        else:
+            pretrained_dict = checkpoint
+
+        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
+        print('missing_keys: ', missing_keys)
+        print('unexpected_keys: ', unexpected_keys)
+        
+    if args.attn_only:
+        for name_p,p in model.named_parameters():
+            if '.attn.' in name_p:
+                p.requires_grad = True
+            else:
+                p.requires_grad = False
+        try:
+            model.head.weight.requires_grad = True
+            model.head.bias.requires_grad = True
+        except:
+            model.fc.weight.requires_grad = True
+            model.fc.bias.requires_grad = True
+        try:
+            model.pos_embed.requires_grad = True
+        except:
+            print('no position encoding')
+        try:
+            for p in model.patch_embed.parameters():
+                p.requires_grad = False
+        except:
+            print('no patch embed')
+            
+    model.to(device)
+
+    model_ema = None
+    if args.model_ema:
+        # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
+        model_ema = ModelEma(
+            model.student.backbone,
+            decay=args.model_ema_decay,
+            device='cpu' if args.model_ema_force_cpu else '',
+            resume='')
+
+    model_without_ddp = model
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
+        model_without_ddp = model.module
+    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print('number of params:', n_parameters)
+    
+    if not args.unscale_lr:
+        linear_scaled_lr = args.lr * args.batch_size * utils.get_world_size() / 512.0
+        args.lr = linear_scaled_lr
+    
+    optimizer = create_optimizer(args, model_without_ddp)
+    loss_scaler = NativeScaler()
+
+    lr_scheduler, _ = create_scheduler(args, optimizer)
+
+    output_dir = Path(args.output_dir)
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume, map_location='cpu')
+
+        model_without_ddp.load_state_dict(checkpoint['model'])
+        if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if args.model_ema:
+                utils._load_checkpoint_for_ema(model_ema, checkpoint['model_ema'])
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+        lr_scheduler.step(args.start_epoch)
+
+    print(f"Start training for {args.epochs} epochs")
+    start_time = time.time()
+    max_accuracy = 0.0
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            data_loader_train.sampler.set_epoch(epoch)
+        
+        train_stats = train_one_epoch(
+            model, data_loader_train,
+            optimizer, device, epoch, loss_scaler,
+            args.clip_grad, model_ema, mixup_fn,
+            set_training_mode=args.train_mode,  # keep in eval mode for deit finetuning / train mode for training and deit III finetuning
+            args = args,
+        )
+
+        lr_scheduler.step(epoch)
+        if args.output_dir:
+            checkpoint_paths = [output_dir / 'checkpoint.pth']
+            for checkpoint_path in checkpoint_paths:
+                utils.save_on_master({
+                    'model': model_without_ddp.state_dict(),
+                    'optimizer': optimizer.state_dict(),
+                    'lr_scheduler': lr_scheduler.state_dict(),
+                    'epoch': epoch,
+                    'model_ema': get_state_dict(model_ema),
+                    'scaler': loss_scaler.state_dict(),
+                    'args': args,
+                }, checkpoint_path)
+
+        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+                     'epoch': epoch,
+                     'n_parameters': n_parameters}
+        
+        if args.output_dir and utils.is_main_process():
+            with (output_dir / "log.txt").open("a") as f:
+                f.write(json.dumps(log_stats) + "\n")
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Training time {}'.format(total_time_str))
+
+
+def train_one_epoch(model: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
+                    model_ema: Optional[ModelEma] = None, mixup_fn: Optional[Mixup] = None,
+                    set_training_mode=True, args = None):
+    model.train(set_training_mode)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 10
+    
+    loader_len = len(data_loader)
+
+    for data_iter_step, inputs_dict in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+
+        for k, v in inputs_dict.items():
+            if isinstance(v, torch.Tensor):
+                inputs_dict[k] = v.to(device, non_blocking=True)
+        
+        with torch.cuda.amp.autocast():
+            loss_dict = model(inputs_dict)
+        
+        loss = loss_dict["loss"]
+        patch_loss = loss_dict["patch_loss"]
+        fea_loss = loss_dict["fea_loss"]
+        token_loss = loss_dict["token_loss"]
+
+        patch_loss_value = patch_loss.item()
+        token_loss_value = token_loss.item()
+        fea_loss_value = fea_loss.item()
+        loss_value = loss.item()
+        
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        optimizer.zero_grad()
+
+        # this attribute is added by timm on one optimizer (adahessian)
+        is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
+        loss_scaler(loss, optimizer, clip_grad=max_norm,
+                    parameters=model.parameters(), create_graph=is_second_order)
+
+        torch.cuda.synchronize()
+        if model_ema is not None:
+            model_ema.update(model.module.student.backbone)
+
+        metric_logger.update(loss=loss_value)
+        metric_logger.update(patch_loss=patch_loss_value)
+        metric_logger.update(token_loss=token_loss_value)
+        metric_logger.update(fea_loss=fea_loss_value)
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('DeiT training and evaluation script', parents=[get_args_parser()])
+    args = parser.parse_args()
+    if args.output_dir:
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    main(args)

1_feature_extractor/models_IB.py

@@ -0,0 +1,40 @@
+import torch.nn as nn
+from compressai.entropy_models import EntropyBottleneck
+from timm.models.vision_transformer import Block
+
+class IF_Module(nn.Module):
+    def __init__(self, embed_dim, num_heads, mlp_ratio, depth=4, norm_layer=nn.LayerNorm):
+        super(IF_Module, self).__init__()
+
+        self.encoder_blocks = nn.ModuleList([   # 4 layers, embed_dim=768, num_heads=12
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for i in range(depth)])
+        self.encoder_norm = norm_layer(embed_dim)
+
+        self.decoder_blocks = nn.ModuleList([   # 4 layers, embed_dim=768, num_heads=12
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for i in range(depth)])
+        
+        self.decoder_norm = norm_layer(embed_dim)
+        self.entropy_bottleneck = EntropyBottleneck(embed_dim)
+
+    def forward(self, x, is_training=False):
+        # ViT analysis transform
+        for blk in self.encoder_blocks:
+            x = blk(x)
+        x = self.encoder_norm(x)
+
+        if is_training:
+            x = x.permute(0, 2, 1)
+            x_hat, x_likelihood = self.entropy_bottleneck(x)
+            x_hat = x_hat.permute(0, 2, 1)
+        else:
+            x_hat = x
+            x_likelihood = None
+
+        # ViT synthesis transform
+        for blk in self.decoder_blocks:
+            x_hat = blk(x_hat)
+        x_hat = self.decoder_norm(x_hat)
+
+        return x_hat, x_likelihood

1_feature_extractor/models_clip.py

@@ -0,0 +1,438 @@
+from collections import OrderedDict
+from typing import Tuple, Union, Callable
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.init import trunc_normal_
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x[:1], key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x.squeeze(0)
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            x = self.relu1(self.bn1(self.conv1(x)))
+            x = self.relu2(self.bn2(self.conv2(x)))
+            x = self.relu3(self.bn3(self.conv3(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.resblocks(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        return x
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.mask_token = nn.Parameter(torch.zeros(1, width))
+
+        self.ln_post = LayerNorm(width)
+
+        self.embed_dim = width
+        self.patch_size = patch_size
+
+        self.init_weights()
+
+    def init_weights(self):
+        trunc_normal_(self.positional_embedding, std=0.02)
+        nn.init.normal_(self.class_embedding, std=1e-6)
+        named_apply(init_weights_vit_timm, self)
+
+    def prepare_tokens_with_masks(self, x, masks=None):    
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        return x
+
+    def forward_features_list(self, x_list, masks_list):
+        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+
+        all_x = [self.transformer(t) for t in x]
+
+        output = []
+        for x, masks in zip(all_x, masks_list):
+            output.append(
+                {
+                    "x_norm_clstoken": self.ln_post(x[:, 0]),
+                    "x_norm_patchtokens": x[:, 1 :],
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward(self, x: torch.Tensor, masks=None):
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self.prepare_tokens_with_masks(x, masks)
+
+        x = self.transformer(x)
+
+        return {
+            "x_norm_clstoken": self.ln_post(x[:, 0]),
+            "x_norm_patchtokens": x[:, 1 :],
+            "x_prenorm": x,
+            "masks": masks,
+        }
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):
+    """ViT weight initialization, original timm impl (for reproducibility)"""
+    if isinstance(module, nn.Linear):
+        trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+
+def vit_small(patch_size=14, teacher_path=None):
+    model = VisionTransformer(
+                input_resolution=224,
+                patch_size=patch_size,
+                width=384,
+                layers=12,
+                heads=6
+            )
+    
+    if teacher_path is not None:
+        checkpoint = torch.load(teacher_path, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            pretrained_dict = checkpoint['state_dict']
+        elif 'model' in checkpoint:
+            pretrained_dict = checkpoint['model']
+        else:
+            pretrained_dict = checkpoint
+
+        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
+        print('missing_keys: ', missing_keys)
+        print('unexpected_keys: ', unexpected_keys)
+    
+    return model
+
+
+def vit_base(patch_size=14, teacher_path=None):
+    model = VisionTransformer(
+                input_resolution=224,
+                patch_size=patch_size,
+                width=768,
+                layers=12,
+                heads=12
+            )
+    
+    if teacher_path is not None:
+        checkpoint = torch.load(teacher_path, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            pretrained_dict = checkpoint['state_dict']
+        elif 'model' in checkpoint:
+            pretrained_dict = checkpoint['model']
+        else:
+            pretrained_dict = checkpoint
+
+        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
+        print('missing_keys: ', missing_keys)
+        print('unexpected_keys: ', unexpected_keys)
+    
+    return model
+
+
+def vit_large(patch_size=14, teacher_path=None):
+    model = VisionTransformer(
+                input_resolution=224,
+                patch_size=patch_size,
+                width=1024,
+                layers=24,
+                heads=16
+            )
+
+    if teacher_path is not None:
+        checkpoint = torch.load(teacher_path, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            pretrained_dict = checkpoint['state_dict']
+        elif 'model' in checkpoint:
+            pretrained_dict = checkpoint['model']
+        else:
+            pretrained_dict = checkpoint
+
+        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
+        print('missing_keys: ', missing_keys)
+        print('unexpected_keys: ', unexpected_keys)
+
+    return model
+
+
+
+if __name__ == "__main__":
+    import argparse
+    import clip
+    import open_clip
+    from fvcore.nn import FlopCountAnalysis, parameter_count_table
+    parser = argparse.ArgumentParser(description='PyTorch resnet Training')
+    args = parser.parse_args()
+
+    # with torch.no_grad():
+        # print(clip.available_models())
+        
+        # device = "cuda" if torch.cuda.is_available() else "cpu"
+        # model, preprocess = clip.load('ViT-L/14', device)
+        # print(model.visual)
+
+        # model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-16', pretrained='laion400m_e32')
+        # model = model.to('cuda')
+
+        # for k,v in model.visual.named_parameters():
+        #     print(k, v.shape)
+
+        # self_model = VisionTransformer(
+        #         input_resolution=224,
+        #         patch_size=32,
+        #         width=768,
+        #         layers=12,
+        #         heads=12
+        #     )
+
+        # print(self_model)
+
+        # for k,v in self_model.named_parameters():
+        #     print(k, v.shape)
+
+        # new_ckpt = OrderedDict()
+        # for k,v in model.visual.named_parameters():
+        #     if 'proj' != k:
+        #         print(k)
+        #         new_ckpt[k] = v
+            # new_ckpt[k] = v
+
+        # torch.save(new_ckpt, '/home/qw/yitian/TA-KD/clip_model/clip_l_14_400m.pth')
+
+        # model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion400m_e32')
+        # print(model.visual)
+
+
+        # model = clip_base_32()
+        # model = clip_base_14()
+
+        # print(parameter_count_table(model))
+
+        # tensor = torch.rand(1, 3, 224, 224)
+        # flops = FlopCountAnalysis(model, tensor)
+        # print("FLOPs: ", flops.total()/1e9)

1_feature_extractor/models_dinov2.py

@@ -0,0 +1,907 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+import math
+import logging
+from typing import Sequence, Tuple, Union, Callable, Dict, Optional, Any, List
+from torch import Tensor
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from torch.nn.init import trunc_normal_
+import os
+import warnings
+
+logger = logging.getLogger("dinov2")
+
+
+from xformers.ops import memory_efficient_attention, unbind, fmha, scaled_index_add, index_select_cat
+from xformers.ops import SwiGLU
+
+XFORMERS_AVAILABLE = True
+warnings.warn("xFormers is available (Attention)")
+
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+    
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = 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, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        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] * self.scale, qkv[1], qkv[2]
+        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 MemEffAttention(Attention):
+    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            if attn_bias is not None:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+
+class SwiGLUFFNFused(SwiGLU):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+        super().__init__(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            out_features=out_features,
+            bias=bias,
+        )
+
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+    Args:
+        img_size: Image size.
+        patch_size: Patch token size.
+        in_chans: Number of input image channels.
+        embed_dim: Number of linear projection output channels.
+        norm_layer: Normalization layer.
+    """
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        _, _, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+
+        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+        x = self.proj(x)  # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)  # B HW C
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
+        return x
+
+    def flops(self) -> float:
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+    
+
+
+
+class Basic_Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = Attention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+    ) -> None:
+        super().__init__()
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    def forward(self, x: Tensor) -> Tensor:
+        def attn_residual_func(x: Tensor) -> Tensor:
+            return self.ls1(self.attn(self.norm1(x)))
+
+        def ffn_residual_func(x: Tensor) -> Tensor:
+            return self.ls2(self.mlp(self.norm2(x)))
+
+        if self.training and self.sample_drop_ratio > 0.1:
+            # the overhead is compensated only for a drop path rate larger than 0.1
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+        elif self.training and self.sample_drop_ratio > 0.0:
+            x = x + self.drop_path1(attn_residual_func(x))
+            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
+        else:
+            x = x + attn_residual_func(x)
+            x = x + ffn_residual_func(x)
+        return x
+
+
+def drop_add_residual_stochastic_depth(
+    x: Tensor,
+    residual_func: Callable[[Tensor], Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> Tensor:
+    # 1) extract subset using permutation
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    x_subset = x[brange]
+
+    # 2) apply residual_func to get residual
+    residual = residual_func(x_subset)
+
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+
+    residual_scale_factor = b / sample_subset_size
+
+    # 3) add the residual
+    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    residual_scale_factor = b / sample_subset_size
+    return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+    if scaling_vector is None:
+        x_flat = x.flatten(1)
+        residual = residual.flatten(1)
+        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    else:
+        x_plus_residual = scaled_index_add(
+            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+        )
+    return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+    """
+    this will perform the index select, cat the tensors, and provide the attn_bias from cache
+    """
+    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+    if all_shapes not in attn_bias_cache.keys():
+        seqlens = []
+        for b, x in zip(batch_sizes, x_list):
+            for _ in range(b):
+                seqlens.append(x.shape[1])
+        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+        attn_bias._batch_sizes = batch_sizes
+        attn_bias_cache[all_shapes] = attn_bias
+
+    if branges is not None:
+        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+    else:
+        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+        cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+    return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[Tensor],
+    residual_func: Callable[[Tensor, Any], Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> Tensor:
+    # 1) generate random set of indices for dropping samples in the batch
+    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+    branges = [s[0] for s in branges_scales]
+    residual_scale_factors = [s[1] for s in branges_scales]
+
+    # 2) get attention bias and index+concat the tensors
+    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+    # 3) apply residual_func to get residual, and split the result
+    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+    outputs = []
+    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+    return outputs
+
+
+class Block(Basic_Block):
+    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
+        """
+        x_list contains a list of tensors to nest together and run
+        """
+        assert isinstance(self.attn, MemEffAttention)
+
+        if self.training and self.sample_drop_ratio > 0.0:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.mlp(self.norm2(x))
+
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            return x_list
+        else:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls2(self.mlp(self.norm2(x)))
+
+            attn_bias, x = get_attn_bias_and_cat(x_list)
+            x = x + attn_residual_func(x, attn_bias=attn_bias)
+            x = x + ffn_residual_func(x)
+            return attn_bias.split(x)
+
+    def forward(self, x_or_x_list):
+        if isinstance(x_or_x_list, Tensor):
+            return super().forward(x_or_x_list)
+        elif isinstance(x_or_x_list, list):
+            if not XFORMERS_AVAILABLE:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise AssertionError    
+
+
+
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+class BlockChunk(nn.ModuleList):
+    def forward(self, x):
+        for b in self:
+            x = b(x)
+        return x
+
+
+class DinoVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        ffn_bias=True,
+        proj_bias=True,
+        drop_path_rate=0.0,
+        drop_path_uniform=False,
+        init_values=None,  # for layerscale: None or 0 => no layerscale
+        embed_layer=PatchEmbed,
+        act_layer=nn.GELU,
+        block_fn=Block,
+        ffn_layer="mlp",
+        block_chunks=1,
+        num_register_tokens=0,
+        interpolate_antialias=False,
+        interpolate_offset=0.1,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            proj_bias (bool): enable bias for proj in attn if True
+            ffn_bias (bool): enable bias for ffn if True
+            drop_path_rate (float): stochastic depth rate
+            drop_path_uniform (bool): apply uniform drop rate across blocks
+            weight_init (str): weight init scheme
+            init_values (float): layer-scale init values
+            embed_layer (nn.Module): patch embedding layer
+            act_layer (nn.Module): MLP activation layer
+            block_fn (nn.Module): transformer block class
+            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+            num_register_tokens: (int) number of extra cls tokens (so-called "registers")
+            interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
+            interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
+        """
+        super().__init__()
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_tokens = 1
+        self.n_blocks = depth
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.interpolate_antialias = interpolate_antialias
+        self.interpolate_offset = interpolate_offset
+
+        self.patch_embed = embed_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 + self.num_tokens, embed_dim))
+        assert num_register_tokens >= 0
+        self.register_tokens = (
+            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
+        )
+
+        if drop_path_uniform is True:
+            dpr = [drop_path_rate] * depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        if ffn_layer == "mlp":
+            logger.info("using MLP layer as FFN")
+            ffn_layer = Mlp
+        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+            logger.info("using SwiGLU layer as FFN")
+            ffn_layer = SwiGLUFFNFused
+        elif ffn_layer == "identity":
+            logger.info("using Identity layer as FFN")
+
+            def f(*args, **kwargs):
+                return nn.Identity()
+
+            ffn_layer = f
+        else:
+            raise NotImplementedError
+
+        blocks_list = [
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                ffn_layer=ffn_layer,
+                init_values=init_values,
+            )
+            for i in range(depth)
+        ]
+        if block_chunks > 0:
+            self.chunked_blocks = True
+            chunked_blocks = []
+            chunksize = depth // block_chunks
+            for i in range(0, depth, chunksize):
+                # this is to keep the block index consistent if we chunk the block list
+                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+        else:
+            self.chunked_blocks = False
+            self.blocks = nn.ModuleList(blocks_list)
+
+        self.norm = norm_layer(embed_dim)
+        self.head = nn.Identity()
+
+        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+        self.init_weights()
+
+    def init_weights(self):
+        trunc_normal_(self.pos_embed, std=0.02)
+        nn.init.normal_(self.cls_token, std=1e-6)
+        if self.register_tokens is not None:
+            nn.init.normal_(self.register_tokens, std=1e-6)
+        named_apply(init_weights_vit_timm, self)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        previous_dtype = x.dtype
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        pos_embed = self.pos_embed.float()
+        class_pos_embed = pos_embed[:, 0]
+        patch_pos_embed = pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_size
+        h0 = h // self.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
+
+        sqrt_N = math.sqrt(N)
+        sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
+            scale_factor=(sx, sy),
+            mode="bicubic",
+            antialias=self.interpolate_antialias,
+        )
+
+        assert int(w0) == patch_pos_embed.shape[-2]
+        assert int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+    def prepare_tokens_with_masks(self, x, masks=None):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+
+        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        if self.register_tokens is not None:
+            x = torch.cat(
+                (
+                    x[:, :1],
+                    self.register_tokens.expand(x.shape[0], -1, -1),
+                    x[:, 1:],
+                ),
+                dim=1,
+            )
+
+        return x
+
+    def forward_features_list(self, x_list, masks_list):
+        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+        for blk in self.blocks:
+            x = blk(x)
+
+        all_x = x
+        output = []
+        for x, masks in zip(all_x, masks_list):
+            x_norm = self.norm(x)
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm[:, 0],
+                    "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward_features(self, x, masks=None):
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self.prepare_tokens_with_masks(x, masks)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x_norm = self.norm(x)
+        return {
+            "x_norm_clstoken": x_norm[:, 0],
+            "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+            "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+            "x_prenorm": x,
+            "masks": masks,
+        }
+
+    def _get_intermediate_layers_not_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        # If n is an int, take the n last blocks. If it's a list, take them
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in blocks_to_take:
+                output.append(x)
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def _get_intermediate_layers_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        output, i, total_block_len = [], 0, len(self.blocks[-1])
+        # If n is an int, take the n last blocks. If it's a list, take them
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for block_chunk in self.blocks:
+            for blk in block_chunk[i:]:  # Passing the nn.Identity()
+                x = blk(x)
+                if i in blocks_to_take:
+                    output.append(x)
+                i += 1
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
+        reshape: bool = False,
+        return_class_token: bool = False,
+        norm=True,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        if self.chunked_blocks:
+            outputs = self._get_intermediate_layers_chunked(x, n)
+        else:
+            outputs = self._get_intermediate_layers_not_chunked(x, n)
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+        class_tokens = [out[:, 0] for out in outputs]
+        outputs = [out[:, 1:] for out in outputs]
+        if reshape:
+            B, _, w, h = x.shape
+            outputs = [
+                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+                for out in outputs
+            ]
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+        return tuple(outputs)
+
+    def forward(self, *args, is_training=False, **kwargs):
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return self.head(ret["x_norm_clstoken"])
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):
+    """ViT weight initialization, original timm impl (for reproducibility)"""
+    if isinstance(module, nn.Linear):
+        trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+def vit_tiny(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=192,
+        depth=12,
+        num_heads=3,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+
+def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
+    """
+    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1536,
+        depth=40,
+        num_heads=24,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+if __name__ == "__main__":
+    import argparse
+    from fvcore.nn import FlopCountAnalysis, parameter_count_table
+
+    with torch.no_grad():
+        model = vit_base(img_size=224,
+        patch_size=14,
+        init_values=1.0,
+        ffn_layer='mlp',
+        block_chunks=0,
+        num_register_tokens=0,
+        interpolate_antialias=False,
+        interpolate_offset=0.1)
+
+        for name, param in model.named_parameters():
+            print(name, param)
+
+        # print(parameter_count_table(model))
+
+        # tensor = torch.rand(1, 3, 224, 224)
+        # flops = FlopCountAnalysis(model, tensor)
+        # print("FLOPs: ", flops.total()/1e9)

1_feature_extractor/models_proteus_clip.py

@@ -0,0 +1,101 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.distributed.nn
+import torch.distributed as dist
+from torch.nn.init import trunc_normal_
+from torch.nn.utils import weight_norm
+import models_clip
+
+
+
+class MetaArch(nn.Module):
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+
+        student_model_dict = dict()
+        teacher_model_dict = dict()
+
+        import_student = getattr(models_clip, cfg.target_model)
+        student = import_student()
+        
+        embed_dim = student.embed_dim
+        
+        import_teacher = getattr(models_clip, cfg.teacher_model)
+        teacher_backbone = import_teacher(teacher_path=cfg.teacher_path)
+        teacher_backbone.eval()
+
+        student_model_dict['backbone'] = student
+        teacher_model_dict['backbone'] = teacher_backbone
+        
+        self.embed_dim = embed_dim
+
+        # initialize parameters and checks
+        self.total_n_global_crops = cfg.batch_size
+
+        self.student = nn.ModuleDict(student_model_dict)
+        self.teacher = nn.ModuleDict(teacher_model_dict)
+
+        teacher_embed_dim = teacher_backbone.embed_dim
+        
+        self.token_head = nn.Sequential(
+                  nn.LayerNorm(embed_dim),
+                  nn.Linear(embed_dim, teacher_embed_dim))
+
+        self.soft_criterion = torch.nn.MSELoss()
+
+        for param in self.teacher.backbone.parameters():
+            param.requires_grad = False
+
+    ## we explicitly remove the patch and feature learning objectives for CLIP training following the original design
+    def forward(self, inputs):
+        global_crops = inputs["collated_global_crops"]
+
+        # compute teacher output
+        # @torch.no_grad()
+        def compute_teacher_output():
+            with torch.no_grad():
+                teacher_backbone_output_dict = self.teacher.backbone(global_crops)
+            teacher_cls_tokens = teacher_backbone_output_dict["x_norm_clstoken"]
+
+            return teacher_cls_tokens
+
+        # get the teacher outputs
+        teacher_cls_tokens = compute_teacher_output()
+
+        student_backbone_output_dict_unmask = self.student.backbone(global_crops)
+
+        student_cls_token_unmask = student_backbone_output_dict_unmask["x_norm_clstoken"]
+
+        ## projection head
+        student_cls_token_unmask = self.token_head(student_cls_token_unmask)
+
+        ## token objective
+        distillation_loss_token = self.soft_criterion(student_cls_token_unmask, teacher_cls_tokens)
+        
+        # coefficient
+        token_loss = self.cfg.lambda_token * distillation_loss_token
+
+        # compute the total loss
+        total_loss = token_loss
+
+        # return the final loss dict
+        loss_dict = {"patch_loss": torch.tensor(0.0), "fea_loss": torch.tensor(0.0), "token_loss": token_loss, "loss": total_loss}
+        
+        return loss_dict

1_feature_extractor/models_proteus_dinov2.py

@@ -0,0 +1,200 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.distributed.nn
+import torch.distributed as dist
+from torch.nn.init import trunc_normal_
+from torch.nn.utils import weight_norm
+import models_dinov2
+from models_IB import IF_Module
+import math
+
+
+class MetaArch(nn.Module):
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+
+        student_model_dict = dict()
+        teacher_model_dict = dict()
+
+        import_student = getattr(models_dinov2, cfg.target_model)
+        student = import_student(img_size=224,
+            patch_size=cfg.patch_size,
+            init_values=1.0,
+            ffn_layer='mlp',
+            block_chunks=0,
+            num_register_tokens=0,
+            interpolate_antialias=False,
+            interpolate_offset=0.1)
+
+        embed_dim = student.embed_dim
+        
+        if cfg.teacher_model == 'vit_base':
+            teacher_backbone = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_lc')
+        elif cfg.teacher_model == 'vit_small':
+            teacher_backbone = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14_lc')
+        elif cfg.teacher_model == 'vit_large':
+            teacher_backbone = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_lc')
+        elif cfg.teacher_model == 'vit_giant':
+            teacher_backbone = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14_lc')
+        teacher_backbone.eval()
+
+        student_model_dict['backbone'] = student
+        teacher_model_dict['backbone'] = teacher_backbone.backbone
+        
+        self.embed_dim = embed_dim
+
+        # initialize parameters and checks
+        self.total_n_global_crops = cfg.batch_size
+
+        self.student = nn.ModuleDict(student_model_dict)
+        self.teacher = nn.ModuleDict(teacher_model_dict)
+
+        teacher_embed_dim = teacher_backbone.backbone.embed_dim
+        self.ibot_head = nn.Sequential(
+                  nn.LayerNorm(embed_dim),
+                  nn.Linear(embed_dim, teacher_embed_dim))
+        
+        self.token_head = nn.Sequential(
+                  nn.LayerNorm(embed_dim),
+                  nn.Linear(embed_dim, teacher_embed_dim))
+
+        self.fea_head = nn.Sequential(
+                  nn.LayerNorm(embed_dim),
+                  nn.Linear(embed_dim, teacher_embed_dim))
+
+        self.soft_criterion = torch.nn.MSELoss()
+
+        self.info_bottleneck = IF_Module(embed_dim=embed_dim, num_heads=12, mlp_ratio=4, depth=4)
+
+        for param in self.teacher.backbone.parameters():
+            param.requires_grad = False
+    
+    def cal_bpp(self, image, unmask_likelihood, mask_likelihood):
+        b, _, h, w = image.size()
+        num_pixels = b * h * w
+        log_unmask_likelihoods = torch.log(unmask_likelihood)
+        log_mask_likelihoods = torch.log(mask_likelihood)
+        bpp = (log_unmask_likelihoods.sum() + log_mask_likelihoods.sum()) / (-math.log(2) * num_pixels * 1.5)
+        return bpp
+
+    def forward(self, inputs):
+        global_crops = inputs["collated_global_crops"]
+        
+        masks = inputs["collated_masks"]
+        mask_indices_list = inputs["mask_indices_list"]
+        n_masked_patches = mask_indices_list.shape[0]
+        upperbound = inputs["upperbound"]
+
+        n_global_crops = 1
+
+        # compute teacher output
+        # @torch.no_grad()
+        def compute_teacher_output():
+            with torch.no_grad():
+                teacher_backbone_output_dict = self.teacher.backbone(global_crops, is_training=True)
+            teacher_cls_tokens = teacher_backbone_output_dict["x_norm_clstoken"]
+            teacher_patch_tokens = teacher_backbone_output_dict["x_norm_patchtokens"]
+            _dim = teacher_patch_tokens.shape[-1]
+
+            # mask teacher patch tokens
+            buffer_tensor_teacher = teacher_patch_tokens.new_zeros(upperbound, _dim)
+            torch.index_select(
+                teacher_patch_tokens.flatten(0, 1),
+                dim=0,
+                index=mask_indices_list,
+                out=buffer_tensor_teacher[:n_masked_patches],
+            )
+            teacher_patch_tokens_masked = buffer_tensor_teacher[:n_masked_patches]
+
+            return teacher_cls_tokens, teacher_patch_tokens, teacher_patch_tokens_masked
+
+        # get the teacher outputs
+        (
+            teacher_cls_tokens,
+            teacher_patch_tokens,
+            teacher_patch_tokens_masked
+        ) = compute_teacher_output()
+        
+        cur_masks = masks if self.cfg.mask_probability > 0 else None
+
+        student_backbone_output_dict, student_backbone_output_dict_unmask = self.student.backbone(
+            [global_crops, global_crops], masks=[cur_masks, None], is_training=True
+        )
+
+        student_cls_token_unmask = student_backbone_output_dict_unmask["x_norm_clstoken"]
+        student_patch_tokens_unmask = student_backbone_output_dict_unmask["x_norm_patchtokens"]
+        student_patch_tokens = student_backbone_output_dict["x_norm_patchtokens"]
+
+        # calculate bitrate
+        student_patch_tokens_unmask, unmask_likelihood = self.info_bottleneck(student_patch_tokens_unmask, is_training=True)
+        student_patch_tokens, mask_likelihood = self.info_bottleneck(student_patch_tokens, is_training=True)
+        bpp = self.cal_bpp(global_crops, unmask_likelihood, mask_likelihood)
+
+        # mask student patch tokens
+        _dim = student_patch_tokens.shape[-1]
+        
+        buffer_tensor_student = student_patch_tokens.new_zeros(upperbound, _dim)
+        buffer_tensor_student[:n_masked_patches].copy_(
+            torch.index_select(student_patch_tokens.flatten(0, 1),
+                                dim=0,
+                                index=mask_indices_list)
+        )
+
+        ## projection head
+        student_patch_tokens_unmask = self.fea_head(student_patch_tokens_unmask)
+        
+        student_cls_token_unmask = self.token_head(student_cls_token_unmask)
+        
+        tokens_after_head = self.ibot_head(buffer_tensor_student)
+        student_patch_tokens_masked = tokens_after_head[:n_masked_patches]
+
+        ## token objective
+        distillation_loss_token = self.soft_criterion(student_cls_token_unmask, teacher_cls_tokens)
+
+        ## fea objective
+        student_whole_fea = torch.cat((student_cls_token_unmask.unsqueeze(1),student_patch_tokens_unmask),dim=1)
+        teacher_whole_fea = torch.cat((teacher_cls_tokens.unsqueeze(1),teacher_patch_tokens),dim=1)
+        distillation_loss_fea = self.soft_criterion(student_whole_fea, teacher_whole_fea)
+
+        ## patch objective
+        patch_loss = self.soft_criterion(student_patch_tokens_masked, teacher_patch_tokens_masked)
+        
+        # coefficient
+        token_loss = self.cfg.lambda_token * distillation_loss_token
+        fea_loss = self.cfg.lambda_fea * distillation_loss_fea
+        patch_loss_weighted = self.cfg.lambda_patch * patch_loss
+        # print(f"self.cfg: {self.cfg}")
+        # print(f"self.cfg.lambda_token: {self.cfg.lambda_token}, self.cfg.lambda_fea: {self.cfg.lambda_fea}, self.cfg.lambda_patch: {self.cfg.lambda_patch}")
+
+        # compute the total loss
+        total_loss = patch_loss_weighted + fea_loss + token_loss + 0.48 * bpp
+        # task_loss = patch_loss + fea_loss + token_loss
+        task_loss = patch_loss + distillation_loss_fea + distillation_loss_token
+
+        # return the final loss dict
+        loss_dict = {"bpp_loss": bpp, 
+                     "patch_loss": patch_loss, 
+                     "fea_loss": distillation_loss_fea, 
+                     "token_loss": token_loss, 
+                     "loss": total_loss,
+                     "task_loss": task_loss,
+                     }
+        
+        return loss_dict

1_feature_extractor/models_proteus_synclr.py

@@ -0,0 +1,161 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.distributed.nn
+import torch.distributed as dist
+from torch.nn.init import trunc_normal_
+from torch.nn.utils import weight_norm
+import models_synclr
+
+
+
+class MetaArch(nn.Module):
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+
+        student_model_dict = dict()
+        teacher_model_dict = dict()
+
+        import_student = getattr(models_synclr, cfg.target_model)
+        student = import_student(patch_size=cfg.patch_size, num_classes=0, mask_style='ibot')
+        
+        embed_dim = student.embed_dim
+        
+        import_teacher = getattr(models_synclr, cfg.teacher_model)
+        teacher_backbone = import_teacher(patch_size=cfg.patch_size, teacher_path=cfg.teacher_path, num_classes=0, mask_style='ibot')
+        teacher_backbone.eval()
+
+        student_model_dict['backbone'] = student
+        teacher_model_dict['backbone'] = teacher_backbone
+        
+        self.embed_dim = embed_dim
+
+        # initialize parameters and checks
+        self.total_n_global_crops = cfg.batch_size
+
+        self.student = nn.ModuleDict(student_model_dict)
+        self.teacher = nn.ModuleDict(teacher_model_dict)
+
+        teacher_embed_dim = teacher_backbone.embed_dim
+        self.patch_head = nn.Sequential(
+                  nn.LayerNorm(embed_dim),
+                  nn.Linear(embed_dim, teacher_embed_dim))
+        
+        self.token_head = nn.Sequential(
+                  nn.LayerNorm(embed_dim),
+                  nn.Linear(embed_dim, teacher_embed_dim))
+
+        self.fea_head = nn.Sequential(
+                  nn.LayerNorm(embed_dim),
+                  nn.Linear(embed_dim, teacher_embed_dim))
+
+        self.soft_criterion = torch.nn.MSELoss()
+
+        for param in self.teacher.backbone.parameters():
+            param.requires_grad = False
+
+    def forward(self, inputs):
+        global_crops = inputs["collated_global_crops"]
+        
+        masks = inputs["collated_masks"]
+        mask_indices_list = inputs["mask_indices_list"]
+        n_masked_patches = mask_indices_list.shape[0]
+        upperbound = inputs["upperbound"]
+
+        n_global_crops = 1
+
+        # compute teacher output
+        # @torch.no_grad()
+        def compute_teacher_output():
+            with torch.no_grad():
+                teacher_backbone_output_dict = self.teacher.backbone(global_crops, is_training=True)
+            teacher_cls_tokens = teacher_backbone_output_dict["x_norm_clstoken"]
+            teacher_patch_tokens = teacher_backbone_output_dict["x_norm_patchtokens"]
+            _dim = teacher_patch_tokens.shape[-1]
+
+            # mask teacher patch tokens
+            buffer_tensor_teacher = teacher_patch_tokens.new_zeros(upperbound, _dim)
+            torch.index_select(
+                teacher_patch_tokens.flatten(0, 1),
+                dim=0,
+                index=mask_indices_list,
+                out=buffer_tensor_teacher[:n_masked_patches],
+            )
+            teacher_patch_tokens_masked = buffer_tensor_teacher[:n_masked_patches]
+
+            return teacher_cls_tokens, teacher_patch_tokens, teacher_patch_tokens_masked
+
+        # get the teacher outputs
+        (
+            teacher_cls_tokens,
+            teacher_patch_tokens,
+            teacher_patch_tokens_masked
+        ) = compute_teacher_output()
+        
+        cur_masks = masks if self.cfg.mask_probability > 0 else None
+
+        student_backbone_output_dict, student_backbone_output_dict_unmask = self.student.backbone(
+            [global_crops, global_crops], masks=[cur_masks, None], is_training=True
+        )
+
+        student_cls_token_unmask = student_backbone_output_dict_unmask["x_norm_clstoken"]
+        student_patch_tokens_unmask = student_backbone_output_dict_unmask["x_norm_patchtokens"]
+        student_patch_tokens = student_backbone_output_dict["x_norm_patchtokens"]
+
+        # mask student patch tokens
+        _dim = student_patch_tokens.shape[-1]
+        
+        buffer_tensor_student = student_patch_tokens.new_zeros(upperbound, _dim)
+        buffer_tensor_student[:n_masked_patches].copy_(
+            torch.index_select(student_patch_tokens.flatten(0, 1),
+                                dim=0,
+                                index=mask_indices_list)
+        )
+
+        ## projection head
+        student_patch_tokens_unmask = self.fea_head(student_patch_tokens_unmask)
+        
+        student_cls_token_unmask = self.token_head(student_cls_token_unmask)
+        
+        tokens_after_head = self.patch_head(buffer_tensor_student)
+        student_patch_tokens_masked = tokens_after_head[:n_masked_patches]
+
+        ## token objective
+        distillation_loss_token = self.soft_criterion(student_cls_token_unmask, teacher_cls_tokens)
+
+        ## fea objective
+        student_whole_fea = torch.cat((student_cls_token_unmask.unsqueeze(1),student_patch_tokens_unmask),dim=1)
+        teacher_whole_fea = torch.cat((teacher_cls_tokens.unsqueeze(1),teacher_patch_tokens),dim=1)
+        distillation_loss_fea = self.soft_criterion(student_whole_fea, teacher_whole_fea)
+
+        ## patch objective
+        patch_loss = self.soft_criterion(student_patch_tokens_masked, teacher_patch_tokens_masked)
+        
+        # coefficient
+        token_loss = self.cfg.lambda_token * distillation_loss_token
+        fea_loss = self.cfg.lambda_fea * distillation_loss_fea
+        patch_loss = self.cfg.lambda_patch * patch_loss
+
+        # compute the total loss
+        total_loss = patch_loss + fea_loss + token_loss
+
+        # return the final loss dict
+        loss_dict = {"patch_loss": patch_loss, "fea_loss": fea_loss, "token_loss": token_loss, "loss": total_loss}
+        
+        return loss_dict

1_feature_extractor/models_synclr.py

@@ -0,0 +1,500 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from functools import partial
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+from timm.models.layers import trunc_normal_, lecun_normal_, to_2tuple
+from timm.models.vision_transformer import Attention
+from timm.models.layers import Mlp, DropPath
+from timm.models.helpers import named_apply
+
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, ffn_targets=False,
+                 return_layer_targets=False):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        # specify the targets for feature regression
+        self.ffn_targets = ffn_targets
+        self.return_layer_targets = return_layer_targets
+
+    def forward(self, x):
+        if isinstance(x, tuple):
+            x = x[0]
+
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        ffn_out = self.mlp(self.norm2(x))
+        x = x + self.drop_path(ffn_out)
+
+        target = ffn_out if self.ffn_targets else x
+
+        if self.return_layer_targets:
+            return x, target
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ 2D Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer
+
+    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
+        - https://arxiv.org/abs/2010.11929
+
+    Includes distillation token & head support for `DeiT: Data-efficient Image Transformers`
+        - https://arxiv.org/abs/2012.12877
+    """
+
+    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=True, representation_size=None, distilled=False,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None,
+                 act_layer=None, weight_init='', ffn_targets=False, return_layer_targets=False):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            num_classes (int): number of classes for classification head
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
+            distilled (bool): model includes a distillation token and head as in DeiT models
+            drop_rate (float): dropout rate
+            attn_drop_rate (float): attention dropout rate
+            drop_path_rate (float): stochastic depth rate
+            embed_layer (nn.Module): patch embedding layer
+            norm_layer: (nn.Module): normalization layer
+            weight_init: (str): weight init scheme
+            ffn_targets (bool): whether we use ffn output or block end as the feature targets
+            return_layer_targets (bool): whether we return every layer targets
+        """
+        super().__init__()
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_tokens = 2 if distilled else 1
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+        act_layer = act_layer or nn.GELU
+
+        self.patch_embed = embed_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.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        self.ffn_targets = ffn_targets
+        self.return_layer_targets = return_layer_targets
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.Sequential(*[
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate,
+                attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer,
+                ffn_targets=ffn_targets, return_layer_targets=return_layer_targets,
+            )
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        # Representation layer
+        if representation_size and not distilled:
+            self.num_features = representation_size
+            self.pre_logits = nn.Sequential(OrderedDict([
+                ('fc', nn.Linear(embed_dim, representation_size)),
+                ('act', nn.Tanh())
+            ]))
+        else:
+            self.pre_logits = nn.Identity()
+
+        # Classifier head(s)
+        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+        self.head_dist = None
+        if distilled:
+            self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()
+
+        self.init_weights(weight_init)
+
+    def init_weights(self, mode=''):
+        assert mode in ('jax', 'jax_nlhb', 'nlhb', '')
+        head_bias = -math.log(self.num_classes) if 'nlhb' in mode else 0.
+        trunc_normal_(self.pos_embed, std=.02)
+        if self.dist_token is not None:
+            trunc_normal_(self.dist_token, std=.02)
+        if mode.startswith('jax'):
+            # leave cls token as zeros to match jax impl
+            named_apply(partial(_init_vit_weights, head_bias=head_bias, jax_impl=True), self)
+        else:
+            trunc_normal_(self.cls_token, std=.02)
+            self.apply(_init_vit_weights)
+
+    def _init_weights(self, m):
+        # this fn left here for compat with downstream users
+        _init_vit_weights(m)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token', 'dist_token'}
+
+    def get_classifier(self):
+        if self.dist_token is None:
+            return self.head
+        else:
+            return self.head, self.head_dist
+
+    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()
+        if self.num_tokens == 2:
+            self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        if self.dist_token is None:
+            x = torch.cat((cls_token, x), dim=1)
+        else:
+            x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
+        x = self.pos_drop(x + self.pos_embed)
+        x = self.blocks(x)
+        x = self.norm(x)
+        if self.dist_token is None:
+            return self.pre_logits(x[:, 0])
+        else:
+            return x[:, 0], x[:, 1]
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        if self.head_dist is not None:
+            x, x_dist = self.head(x[0]), self.head_dist(x[1])  # x must be a tuple
+            if self.training and not torch.jit.is_scripting():
+                # during inference, return the average of both classifier predictions
+                return x, x_dist
+            else:
+                return (x + x_dist) / 2
+        else:
+            x = self.head(x)
+        return x
+
+
+def _init_vit_weights(module: nn.Module, name: str = '', head_bias: float = 0., jax_impl: bool = False):
+    """ ViT weight initialization
+    * When called without n, head_bias, jax_impl args it will behave exactly the same
+      as my original init for compatibility with prev hparam / downstream use cases (ie DeiT).
+    * When called w/ valid n (module name) and jax_impl=True, will (hopefully) match JAX impl
+    """
+    if isinstance(module, nn.Linear):
+        if name.startswith('head'):
+            nn.init.zeros_(module.weight)
+            nn.init.constant_(module.bias, head_bias)
+        elif name.startswith('pre_logits'):
+            lecun_normal_(module.weight)
+            nn.init.zeros_(module.bias)
+        else:
+            if jax_impl:
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    if 'mlp' in name:
+                        nn.init.normal_(module.bias, std=1e-6)
+                    else:
+                        nn.init.zeros_(module.bias)
+            else:
+                trunc_normal_(module.weight, std=.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+    elif jax_impl and isinstance(module, nn.Conv2d):
+        # NOTE conv was left to pytorch default in my original init
+        lecun_normal_(module.weight)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+    elif isinstance(module, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
+        nn.init.zeros_(module.bias)
+        nn.init.ones_(module.weight)
+
+
+
+def compute_gather_ids(masks):
+    unmask_indices = masks.logical_not().nonzero(as_tuple=False)
+    ids_keep = unmask_indices[:, -1].reshape(masks.shape[0], -1)
+    return ids_keep
+
+
+class MaskedTransformer(VisionTransformer):
+    """Inherit vision transformer from timm"""
+
+    def __init__(self, mask_style='ibot', **kwargs):
+        super().__init__(**kwargs)
+        assert mask_style in ["ibot", "mae", "none"], "mask_style must be `ibot`, `mae`, or `none`"
+
+        self.patch_size = self.patch_embed.patch_size
+        if isinstance(self.patch_size, tuple):
+            self.patch_size = self.patch_size[0]
+
+        nn.init.normal_(self.cls_token, std=1e-6)
+
+        self.mask_style = mask_style
+        if self.mask_style == "ibot":
+            self.mask_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+            torch.nn.init.normal_(self.mask_token, std=.02)
+
+    def interpolate_pos_encoding(self, x, w, h, npatch):
+        previous_dtype = x.dtype
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        pos_embed = self.pos_embed.float()
+        class_pos_embed = pos_embed[:, 0]
+        patch_pos_embed = pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_size
+        h0 = h // self.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode="bicubic",
+        )
+
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+    def prepare_tokens_with_masks(self, x, masks=None):
+        """
+        Args:
+            x: data w/ shape [b, c, h, w]
+            masks: shape [b, n], n is the number of tokens, 1 means masked, 0 means unmasked
+        """
+        b, c, h, w = x.shape
+        x = self.patch_embed(x)
+        if masks is not None:
+            if self.mask_style == 'ibot':
+                x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype), x)
+            elif self.mask_style == 'mae':  # only gather unmasked patches
+                # add pos_embed before shuffle
+                pos_embed = self.interpolate_pos_encoding(x, w, h, npatch=x.shape[1])
+                x = x + pos_embed[:, 1:, :]
+                ids_keep = compute_gather_ids(masks)
+                x = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, x.shape[-1]))
+                # x = x[masks.logical_not()]
+                # x = x.reshape(b, -1, x.size(-1))
+            else:
+                raise NotImplementedError(f"mask style {self.mask_style} is not supported")
+
+        if (masks is None) or (self.mask_style != "mae"):
+            x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+            x = x + self.interpolate_pos_encoding(x, w, h, npatch=x.shape[1]-1)
+        else:
+            # mae-style masking, only need to add cls tokens w/ pos embedding
+            cls_token = self.cls_token + self.pos_embed[:, :1, :]
+            x = torch.cat((cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+
+        return x
+
+    def forward_features_list(self, x_list, masks_list):
+        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+
+        num_data = len(x)
+        if self.return_layer_targets:
+            all_layer_results = [[] for _ in range(num_data)]
+            for i, blk in enumerate(self.blocks):
+                out = [blk(t) for t in x]
+                x = [o[0] for o in out]
+                # store layer targets
+                for j in range(num_data):
+                    all_layer_results[j].append(out[j][1])
+            all_x = x
+        else:
+            all_x = [self.blocks(t) for t in x]
+            all_layer_results = [None for _ in range(num_data)]
+
+        output = []
+        for x, masks, layer_results in zip(all_x, masks_list, all_layer_results):
+            x_norm = self.norm(x)
+            output.append(
+                {
+                    "x_norm": x_norm,
+                    "x_norm_clstoken": x_norm[:, 0],
+                    "x_norm_patchtokens": x_norm[:, 1:],
+                    "masks": masks,
+                    "layer_results": layer_results,
+                }
+            )
+        return output
+
+    def forward_features(self, x, masks=None):
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self.prepare_tokens_with_masks(x, masks)
+
+        if self.return_layer_targets:
+            layer_results = []
+            for i, blk in enumerate(self.blocks):
+                x, lr = blk(x)
+                layer_results.append(lr)
+        else:
+            x = self.blocks(x)
+            layer_results = None
+
+        x_norm = self.norm(x)
+        return {
+            "x_norm": x_norm,
+            "x_norm_clstoken": x_norm[:, 0],
+            "x_norm_patchtokens": x_norm[:, 1:],
+            "masks": masks,
+            "layer_results": layer_results,
+        }
+
+    def forward(self, *args, is_training=False, **kwargs):
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return ret["x_norm_clstoken"]
+
+
+def vit_small(patch_size=16, teacher_path=None, **kwargs):
+    model = MaskedTransformer(
+        patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, **kwargs)
+
+    if teacher_path is not None:
+        checkpoint = torch.load(teacher_path, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            pretrained_dict = checkpoint['state_dict']
+        elif 'model' in checkpoint:
+            pretrained_dict = checkpoint['model']
+        else:
+            pretrained_dict = checkpoint
+
+        pretrained_dict = {k.replace("module.visual.", ""): v for k, v in pretrained_dict.items()}
+
+        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
+        print('missing_keys: ', missing_keys)
+        print('unexpected_keys: ', unexpected_keys)
+    
+    return model
+
+
+def vit_base(patch_size=16, teacher_path=None, **kwargs):
+    model = MaskedTransformer(
+        patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, **kwargs)
+
+    if teacher_path is not None:
+        checkpoint = torch.load(teacher_path, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            pretrained_dict = checkpoint['state_dict']
+        elif 'model' in checkpoint:
+            pretrained_dict = checkpoint['model']
+        else:
+            pretrained_dict = checkpoint
+
+        pretrained_dict = {k.replace("module.visual.", ""): v for k, v in pretrained_dict.items()}
+
+        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
+        print('missing_keys: ', missing_keys)
+        print('unexpected_keys: ', unexpected_keys)
+
+    return model
+
+
+def vit_large(patch_size=14, teacher_path=None, **kwargs):
+    model = MaskedTransformer(
+        patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16, **kwargs)
+    
+    if teacher_path is not None:
+        checkpoint = torch.load(teacher_path, map_location='cpu')
+
+        if 'state_dict' in checkpoint:
+            pretrained_dict = checkpoint['state_dict']
+        elif 'model' in checkpoint:
+            pretrained_dict = checkpoint['model']
+        else:
+            pretrained_dict = checkpoint
+
+        pretrained_dict = {k.replace("module.visual.", ""): v for k, v in pretrained_dict.items()}
+
+        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
+        print('missing_keys: ', missing_keys)
+        print('unexpected_keys: ', unexpected_keys)
+
+    return model
+
+
+if __name__ == '__main__':
+    import argparse
+    from fvcore.nn import FlopCountAnalysis, parameter_count_table
+    parser = argparse.ArgumentParser(description='PyTorch resnet Training')
+    args = parser.parse_args()
+
+    with torch.no_grad():
+        model = vit_base(patch_size=14, num_classes=0, mask_style='ibot')
+        
+        # x = torch.randn(1, 3, 224, 224)
+        # out = model(x)
+        # print(out.shape)
+
+        print(parameter_count_table(model))
+
+        tensor = torch.rand(1, 3, 224, 224)
+        flops = FlopCountAnalysis(model, tensor)
+        print("FLOPs: ", flops.total()/1e9)

1_feature_extractor/requirements.txt

@@ -0,0 +1,134 @@
+absl-py==2.1.0
+accelerate==0.33.0
+aiohttp==3.9.5
+aiosignal==1.3.1
+antlr4-python3-runtime==4.9.3
+async-timeout==4.0.3
+attrs==23.2.0
+cachetools==5.4.0
+certifi==2024.7.4
+charset-normalizer==3.3.2
+click==8.1.7
+cloudpickle==3.0.0
+cmake==3.30.1
+compressai==1.2.6
+contourpy==1.2.1
+cubinlinker-cu11==0.3.0.post2
+cuda-python==11.8.2
+cudf-cu11==24.6.1
+cuml-cu11==24.6.1
+cupy-cuda11x==13.2.0
+cycler==0.12.1
+Cython==3.0.10
+dask==2024.5.1
+dask-cuda==24.6.0
+dask-cudf-cu11==24.6.1
+dask-expr==1.1.1
+diffusers==0.30.0
+distributed==2024.5.1
+distributed-ucxx-cu11==0.38.0
+einops==0.8.0
+fastrlock==0.8.2
+filelock==3.15.4
+fonttools==4.53.1
+frozenlist==1.4.1
+fsspec==2024.6.1
+ftfy==6.2.0
+future==1.0.0
+fvcore==0.1.5.post20221221
+grpcio==1.65.2
+h5py==3.11.0
+huggingface-hub==0.24.5
+idna==3.7
+importlib_metadata==8.0.0
+importlib_resources==6.4.0
+iopath==0.1.10
+Jinja2==3.1.4
+joblib==1.4.2
+kiwisolver==1.4.5
+libucx-cu11==1.15.0.post1
+lightning-utilities==0.11.6
+lit==18.1.8
+llvmlite==0.43.0
+locket==1.0.0
+Markdown==3.6
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.1
+mdurl==0.1.2
+mpmath==1.3.0
+msgpack==1.0.8
+multidict==6.0.5
+mypy-extensions==1.0.0
+networkx==3.2.1
+numba==0.60.0
+numpy==1.26.4
+nvtx==0.2.10
+omegaconf==2.3.0
+open-clip-torch==2.0.2
+opencv-python==4.10.0.84
+packaging==24.1
+pandas==2.2.2
+partd==1.4.2
+pillow==10.4.0
+portalocker==2.10.1
+protobuf==4.25.4
+psutil==6.0.0
+ptxcompiler-cu11==0.8.1.post1
+pyarrow==16.1.0
+pycocotools==2.0.8
+pyDeprecate==0.3.1
+Pygments==2.18.0
+pylibraft-cu11==24.6.0
+pynvml==11.4.1
+pyparsing==3.1.2
+pyre-extensions==0.0.23
+python-dateutil==2.9.0.post0
+pytorch-lightning==1.5.0
+pytorch-msssim==1.0.0
+pytz==2024.1
+PyYAML==6.0.1
+raft-dask-cu11==24.6.0
+rapids-dask-dependency==24.6.0
+regex==2024.7.24
+requests==2.32.3
+rich==13.7.1
+rmm-cu11==24.6.0
+safetensors==0.4.3
+scikit-learn==1.5.1
+scipy==1.13.1
+six==1.16.0
+sortedcontainers==2.4.0
+submitit==1.5.1
+sympy==1.13.1
+tabulate==0.9.0
+tblib==3.0.0
+tensorboard==2.17.0
+tensorboard-data-server==0.7.2
+termcolor==2.4.0
+threadpoolctl==3.5.0
+timm==1.0.7
+tokenizers==0.13.3
+toolz==0.12.1
+torch==2.0.0+cu117
+torch_geometric==2.5.3
+torchmetrics==0.10.3
+torchvision==0.15.0+cu117
+tornado==6.4.1
+tqdm==4.66.4
+transformers==4.27.0
+treelite==4.1.2
+triton==2.0.0
+typing-inspect==0.9.0
+typing_extensions==4.12.2
+tzdata==2024.1
+ucx-py-cu11==0.38.0
+ucxx-cu11==0.38.0
+urllib3==2.2.2
+wcwidth==0.2.13
+Werkzeug==3.0.3
+xformers==0.0.18
+yacs==0.1.8
+yarl==1.9.4
+zict==3.0.0
+zipp==3.19.2