text: # Copyright (c) ByteDance, Inc. and its affiliates.
  # All rights reserved.
  #
  # This source code is licensed under the license found in the
  # LICENSE file in the root directory of this source tree.
  
  import argparse
  import datetime
  import numpy as np
  import os
  import time
  import random
  import torch
  import torch.nn as nn
  import torch.backends.cudnn as cudnn
  import json
  from pathlib import Path
  from objective import *
  from timm.scheduler import create_scheduler
  from timm.utils import get_state_dict
  from datasets import ImageNet, ImageNetLMDB
  from augmentation import get_augmentations
  from engine import train_one_epoch, eval_one_epoch, inference
  import utils
  import warnings
  from tensorboardX import SummaryWriter
  import torchvision
  #from torchvision.models import resnet50
  #from widen_resnet import resnet50w2, resnet50w4, resnet50w5, resnet200, resnet200w2
  import widen_resnet
  from lars import *
  import vision_transformer as vit
  from timm.utils import ModelEmaV2 as ModelEma
  from functools import partial
  from torchvision import transforms
  warnings.filterwarnings('ignore')
  
  
  def get_args_parser():
      parser = argparse.ArgumentParser('Self-Supervised', add_help=False)
      parser.add_argument('--dataset', type=str, default='imagenet')
      parser.add_argument('--lam1', type=float, default=0.0, metavar='LR')
      parser.add_argument('--lam2', type=float, default=1.0, metavar='LR')
      parser.add_argument('--tau', type=float, default=1.0, metavar='LR')
      parser.add_argument('--lbn_type', type=str, default='bn')
      parser.add_argument('--determine', type=int, default=0)
      parser.add_argument('--aug', type=str, default='multicrop')
      parser.add_argument('--img_size', type=int, default=224)
      parser.add_argument('--drop', type=float, default=0.0)
      parser.add_argument('--clip_norm', type=float, default=0.0)
      parser.add_argument('--EPS', type=float, default=1e-5, help='episillon')
      parser.add_argument('--reduce_mean', type=float, default=0)
  
      parser.add_argument('--eval_only', type=int, default=0)
      parser.add_argument('--inference_only', type=int, default=0)
      parser.add_argument('--img_path', type=str, default='./test.jpg')
      parser.add_argument('--match_path', type=str, default='')
  
      parser.add_argument('--threshold', type=float, default=0.0)
      parser.add_argument('--quantile', type=float, default=0.5)
      parser.add_argument('--quantile_end', type=float, default=0.6)
      parser.add_argument('--enable_watch', type=int, default=1)
  
      parser.add_argument('--use_momentum_encoder', type=int, default=0)
      parser.add_argument('--momentum_start', default=0.996, type=float)
      parser.add_argument('--momentum_end', default=1.0, type=float)
      parser.add_argument('--freeze_embedding', default=0, type=int)
  
      # self-label relevant
      parser.add_argument('--mme_epochs', type=int, default=800)
      parser.add_argument('--sl_warmup_epochs', type=int, default=5)
      parser.add_argument('--lr_sl', type=float, default=0.05, metavar='LR')
  
      # maybe different between cnn and vit.
      parser.add_argument('--act', type=str, default='relu')
      parser.add_argument('--patch-size', default=16, type=int)
      parser.add_argument('--backbone', default='resnet50', type=str, metavar='BACKBONEMODEL', help='Name of model to train')
      parser.add_argument('--weight-decay', type=float, default=1.5e-6, help='weight decay (default: 1e-4)')
      parser.add_argument('--weight-decay-end', type=float, default=1.5e-6)
      parser.add_argument('--optim', default='lars', type=str)
      parser.add_argument('--lr', type=float, default=0.5, metavar='LR')
      parser.add_argument('--proj_trunc_init', type=int, default=0)
      parser.add_argument('--proj_norm', type=str, default='bn', choices=['bn', 'ln', 'none'])
      parser.add_argument('--drop_path', type=float, default=0.0)
      
      
      # multi-crop enabled only the aug is set to multicrop
      parser.add_argument('--local_crops_number', type=int, default=12,)
      parser.add_argument('--crops_interact_style', type=str, default='sparse')
      parser.add_argument('--min1', type=float, default=0.4, metavar='LR')
      parser.add_argument('--max1', type=float, default=1.0, metavar='LR')
      parser.add_argument('--min2', type=float, default=0.05, metavar='LR')
      parser.add_argument('--max2', type=float, default=0.4, metavar='LR')
  
  
      parser.add_argument('--batch-size', default=128, type=int)
      parser.add_argument('--bunch-size', default=256, type=int)
  
      parser.add_argument('--epochs', default=850, type=int)
      parser.add_argument('--dim', default=4096, type=int)
      parser.add_argument('--hid_dim', default=4096, type=int)
  
      parser.add_argument('--eval-only', default=0, type=int)
  
      # Model parameters
  
      parser.add_argument('--exclude-bias-weight-decay', type=int, default=1)
  
      # Learning rate schedule parameters
      parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER',
                          help='LR scheduler (default: "cosine"')
      parser.add_argument('--lr_wbr', type=float, default=1.0)
      parser.add_argument('--warmup-epochs', type=int, default=10, metavar='N',
                          help='epochs to warmup LR, if scheduler supports')
  
  
      # Dataset parameters
      parser.add_argument('--data-path', default='/nothing/', type=str, help='dataset path')
      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('--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)
      parser.add_argument('--use-lmdb', action='store_true')
      parser.add_argument('--amp', default=1, type=int)
  
      # distributed training parameters
      parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes')
      parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')
      return parser
  
  class TWIST(nn.Module):
      def __init__(self, args):
          super(TWIST, self).__init__()
          if args.backbone.startswith('resnet'):
              widen_resnet.__dict__['resnet50'] = torchvision.models.resnet50
              self.backbone = widen_resnet.__dict__[args.backbone]()
              self.feature_dim = self.backbone.fc.weight.shape[1]
          else: # vision transformer based models
              if args.backbone.startswith('vit'):
                  self.backbone = vit.__dict__[args.backbone](
                          patch_size=args.patch_size, 
                          norm_layer=(partial(nn.LayerNorm, eps=1e-6)), 
                          drop_path_rate=args.drop_path,
                          freeze_embedding=args.freeze_embedding,
                  )
                  self.feature_dim = self.backbone.embed_dim
          self.backbone.fc = nn.Identity()
          self.backbone = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.backbone)
          if args.crops_interact_style != 'class':
              self.projection_heads = ProjectionHead(args, feature_dim=self.feature_dim)
          else:
              self.projection_heads = utils.ClassHead(args, feature_dim=2048)
  
      def forward(self, x):
          """
              Codes about multi-crop is borrowed from the codes of Dino
              https://github.com/facebookresearch/dino
          """
          if not isinstance(x, list):
              x = [x]
          # the first indices of aug with changing resolution
          idx_crops = torch.cumsum(torch.unique_consecutive(
              torch.tensor([inp.shape[-1] for inp in x]),
              return_counts=True,
          )[1], 0)
  
          start_idx = 0
          for end_idx in idx_crops:
              _out = self.backbone(torch.cat(x[start_idx: end_idx]))
              if start_idx == 0:
                  output = _out
              else:
                  output = torch.cat((output, _out))
              start_idx = end_idx
  
          out = self.projection_heads(output)
          return out
  
      def backbone_weights(self):
          return self.backbone.state_dict()
  
  class ProjectionHead(nn.Module):
      def __init__(self, args, feature_dim=2048):
          super(ProjectionHead, self).__init__()
          if args.lbn_type == 'bn':
              assert args.bunch_size % args.batch_size == 0
              ranks = list(range(utils.get_world_size()))
              print('---ALL RANKS----\n{}'.format(ranks))
              procs_per_bunch = args.bunch_size // args.batch_size
              assert utils.get_world_size() % procs_per_bunch == 0
              n_bunch = utils.get_world_size() // procs_per_bunch
              rank_groups = [ranks[i*procs_per_bunch: (i+1)*procs_per_bunch] for i in range(n_bunch)]
              print('---RANK GROUPS----\n{}'.format(rank_groups))
              process_groups = [torch.distributed.new_group(pids) for pids in rank_groups]
              bunch_id = utils.get_rank() // procs_per_bunch
              process_group = process_groups[bunch_id]
              print('---CURRENT GROUP----\n{}'.format(process_group))
              norm = nn.SyncBatchNorm(args.dim, affine=0, process_group=process_group)
          elif args.lbn_type == 'syncbn':
              norm = nn.SyncBatchNorm(args.dim, affine=0)
          elif args.lbn_type == 'identity':
              norm = nn.Identity()
          else:
              raise NotImplementedError
  
          if args.proj_norm == 'bn':
              batchnorm = nn.SyncBatchNorm
          elif args.proj_norm == 'ln':
              batchnorm = partial(nn.LayerNorm, eps=1e-6)
          elif args.proj_norm == 'none':
              batchnorm = nn.Identity
          else:
              raise NotImplementedError
  
          self.projection_head = nn.Sequential(
              nn.Linear(feature_dim, args.hid_dim, bias=True),
              batchnorm(args.hid_dim),
              nn.ReLU() if args.act == 'relu' else nn.GELU(),
              nn.Dropout(p=args.drop),
  
              nn.Linear(args.hid_dim, args.hid_dim, bias=True),
              batchnorm(args.hid_dim),
              nn.ReLU() if args.act == 'relu' else nn.GELU(),
          )
  
          last_linear = nn.Linear(args.hid_dim, args.dim, bias=True)
          self.last_linear = last_linear
          self.norm = norm
  
          if args.backbone.startswith('vit') or args.proj_trunc_init:
              print('using vit initialization')
              self.apply(self._vit_init_weights)
  
      def _vit_init_weights(self, m):
          if isinstance(m, nn.Linear):
              utils.trunc_normal_(m.weight, std=.02)
              if isinstance(m, nn.Linear) and m.bias is not None:
                  nn.init.constant_(m.bias, 0)
  
      def reg_gnf(self, grad):
          self.gn_f = grad.abs().mean().item()
  
      def reg_gnft(self, grad):
          self.gn_ft = grad.abs().mean().item()
  
      def forward(self, x):
          x = self.projection_head(x)
          f = self.last_linear(x)
          ft = self.norm(f)
          if self.train and x.requires_grad:
              f.register_hook(self.reg_gnf)
              ft.register_hook(self.reg_gnft)
          self.f_column_std = f.std(dim=0, unbiased=False).mean()
          self.f_row_std    = f.std(dim=1, unbiased=False).mean()
          self.ft_column_std = ft.std(dim=0, unbiased=False).mean()
          self.ft_row_std    = ft.std(dim=1, unbiased=False).mean()
  
          return ft
  
  def main(args):
      if args.epochs <=20:
          args.warmup_epochs = 3 
      elif args.epochs <= 50:
          args.warmup_epochs = 5 
      else:
          args.warmup_epochs = 10
      if args.mme_epochs > args.epochs:
          args.mme_epochs = args.epochs
      if args.crops_interact_style == 'self_label':
          assert args.mme_epochs == args.epochs
  
      utils.init_distributed_mode(args)
      print(args)
  
      output_dir = Path(args.output_dir)
      args.global_crops_scale = (args.min1, args.max1)
      args.local_crops_scale = (args.min2, args.max2)
      print('global crops: {}'.format(args.global_crops_scale))
      print('local  crops: {}'.format(args.local_crops_scale))
  
      cudnn.benchmark = True
      device = torch.device(args.device)
      num_tasks = utils.get_world_size()
      global_rank = utils.get_rank()
  
      seed = args.seed + utils.get_rank()
      torch.manual_seed(seed)
      np.random.seed(seed)
      if args.determine:
          seed = args.seed + utils.get_rank()
          torch.manual_seed(seed)
          random.seed(seed)
          torch.cuda.manual_seed(seed)
          torch.cuda.manual_seed_all(seed)
          np.random.seed(seed)
          torch.backends.cudnn.deterministic = True
          torch.backends.cudnn.benchmark = False
  
      # =================== Data Preparation =================== 
      if args.use_lmdb:
          dataset_train = ImageNetLMDB(args.data_path, 'train.lmdb', get_augmentations(args))
      else:
          dataset_train = ImageNet(os.path.join(args.data_path, 'train'), get_augmentations(args))
  
      sampler_train = torch.utils.data.DistributedSampler(
          dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
      )
  
      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,
      )
      
      # prepare evaluation data loader to make unsupervised classification
      if args.dim == 1000 or args.eval_only:
          normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
          val_aug = transforms.Compose([
              transforms.Resize(256),
              transforms.CenterCrop(224),
              transforms.ToTensor(),
              normalize,
          ])
          if args.use_lmdb:
              dataset_val = ImageNetLMDB(args.data_path, 'val.lmdb', val_aug)
          else:
              dataset_val = ImageNet(os.path.join(args.data_path, 'val'), val_aug)
  
          sampler_val = torch.utils.data.DistributedSampler(
              dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False
          )
          data_loader_val = torch.utils.data.DataLoader(
              dataset_val, sampler=sampler_val,
              batch_size=args.batch_size,
              num_workers=args.num_workers,
              pin_memory=args.pin_mem,
              drop_last=False,
          )
  
      # ===================   Model   =================== 
      model = TWIST(args)
      model.to(device)
  
      if args.use_momentum_encoder:
          teacher_model = TWIST(args)
          teacher_model.to(device)
      else:
          teacher_model = None
  
      n_bb_parameters = sum(p.numel() for p in model.backbone.parameters() if p.requires_grad)
      print('number of backbone params:{:.2f} M'.format(n_bb_parameters/1e6))
      n_ph_parameters = sum(p.numel() for p in model.projection_heads.parameters() if p.requires_grad)
      print('number of head params:{:.2f} M'.format(n_ph_parameters/1e6))
      n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
      print('number of total params:{:.2f} M'.format(n_parameters/1e6))
  
      model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
      model_without_ddp = model.module
      if args.use_momentum_encoder:
          teacher_model = torch.nn.parallel.DistributedDataParallel(teacher_model, device_ids=[args.gpu])
          teacher_model_without_ddp = teacher_model.module
          teacher_model_without_ddp.load_state_dict(model_without_ddp.state_dict())
          for p in teacher_model.parameters():
              p.requires_grad = False
          momentum_schedule = utils.cosine_scheduler(args.momentum_start, args.momentum_end, args.mme_epochs, len(data_loader_train))
      else:
          momentum_schedule = None
  
      # =================== Optimizer =================== 
      linear_scaled_lr = args.lr * args.batch_size * utils.get_world_size() / 256.0
      args.lr = linear_scaled_lr
      linear_scaled_lr_sl = args.lr_sl * args.batch_size * utils.get_world_size() / 256.0
      args.lr_sl = linear_scaled_lr_sl
  
      param_weights = []
      param_biases = []
      for name, param in model.named_parameters():
          if not param.requires_grad:
              print('{} is not optimized'.format(name))
              continue
          skip = ['pos_embed', 'cls_token', 'dist_token']
          if len(param.shape) == 1 or name.endswith(".bias") or sum([sk in name for sk in skip]):
              print('{} has been excluded for weight decay'.format(name))
              param_biases.append(param)
          else:
              param_weights.append(param)
      
      if args.optim == 'sgd':
          bias_weight_decay = 0.0 if args.exclude_bias_weight_decay else args.weight_decay
          parameters = [{'params': param_weights, 'weight_decay': args.weight_decay}, 
                        {'params': param_biases,  'weight_decay': bias_weight_decay}]
          optimizer = torch.optim.SGD(parameters, lr=0, momentum=0.9, weight_decay=args.weight_decay)
      elif args.optim in ['lars', 'lars_oss']:
          bias_weight_decay = 0.0 if args.exclude_bias_weight_decay else args.weight_decay
          parameters = [{'params': param_weights, 'weight_decay': args.weight_decay, 'lars_exclude': False}, 
                        {'params': param_biases,  'weight_decay': bias_weight_decay, 'lars_exclude': True}]
          optimizer = LARS_OPENSELF(parameters, lr=0, weight_decay=args.weight_decay, momentum=0.9)
      elif args.optim == 'admw':
          bias_weight_decay = 0.0 if args.exclude_bias_weight_decay else args.weight_decay
          parameters = [{'params': param_weights, 'weight_decay': args.weight_decay}, 
                        {'params': param_biases,  'weight_decay': bias_weight_decay}]
          optimizer = torch.optim.AdamW(parameters)
  
      # weight decay scheduler, used for vision transformer proposed by dino.
      wd_schedule = utils.cosine_scheduler(
          args.weight_decay,
          args.weight_decay_end,
          args.epochs, 
          len(data_loader_train),
      )
      qt_schedule = utils.cosine_scheduler(
          args.quantile,
          args.quantile_end,
          args.epochs if args.crops_interact_style=='self_label' else (args.epochs-args.mme_epochs),
          len(data_loader_train),
      )
  
      # =================== Loss Function =================== 
      criterion = EntLoss(args, args.lam1, args.lam2, pqueue=None)
      loss_scaler = torch.cuda.amp.GradScaler() if args.amp else None
  
      if utils.is_main_process(): # Tensorboard configuration
          local_runs = os.path.join(args.output_dir, 'runs_{}'.format( args.output_dir.replace('/','').replace('.','') ))
          writer = SummaryWriter(logdir=local_runs)
  
      # =================== Traning =================== 
      if args.resume: # Resume traning from checkpoint
          checkpoint = torch.load(args.resume, map_location='cpu')
          try:
              model_without_ddp.load_state_dict(checkpoint['model'])
          except:
              ckpt = checkpoint['model']
              unexpected = {"projection_heads.cls_heads.0.0.weight": "projection_heads.last_linear.weight", 
                            "projection_heads.cls_heads.0.0.bias":  "projection_heads.last_linear.bias",
                            "projection_heads.cls_heads.0.1.running_mean": "projection_heads.norm.running_mean", 
                            "projection_heads.cls_heads.0.1.running_var": "projection_heads.norm.running_var", 
                            "projection_heads.cls_heads.0.1.num_batches_tracked": "projection_heads.norm.num_batches_tracked"
                           }
              for k, v in unexpected.items():
                  ckpt[v] = ckpt[k]
                  del ckpt[k]
              model_without_ddp.load_state_dict(ckpt)
          if 'optimizer' in checkpoint and 'epoch' in checkpoint:
              optimizer.load_state_dict(checkpoint['optimizer'])
              args.start_epoch = checkpoint['epoch'] + 1
              if 'scaler' in checkpoint:
                  loss_scaler.load_state_dict(checkpoint['scaler'])
          if 'teacher' in checkpoint:
              teacher_model_without_ddp.load_state_dict(checkpoint['teacher'])
      
      if args.eval_only:
          eval_stats = eval_one_epoch(args, 
                  model, data_loader_val, device, 
                  logfn=os.path.join(output_dir, 'detail_log.txt')
          )
          np.save(os.path.join(args.output_dir, 'match.npy'), eval_stats['match'])
          np.save(os.path.join(args.output_dir, 'mapped_preds.npy'), eval_stats['mapped_preds'])
          print({k: v for k,v in eval_stats.items() if k not in ['match', 'mapped_preds']})
          return
  
      if args.inference_only:
          inference(args, model, args.img_path, device, np.load(args.match_path))
          return
  
      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 epoch == args.mme_epochs:
              print("Switching to Harder Multi Crop")
              args.global_crops_scale = (0.14, 1.0)
              args.local_crops_scale = (0.05, 0.14)
              args.aug = 'multicrop'
              data_loader_train.dataset.aug = get_augmentations(args)
  
          data_loader_train.sampler.set_epoch(epoch)
          
          # train for one epoch
          train_stats = train_one_epoch(args,
              model, criterion, data_loader_train,
              optimizer, device, epoch, 
              set_training_mode=True,
              scaler=loss_scaler,
              logfn=os.path.join(output_dir, 'detail_log.txt'),
              wd_schedule=wd_schedule,
              qt_schedule=qt_schedule,
              teacher_model=teacher_model,
              momentum_schedule=momentum_schedule,
          )
  
          if args.crops_interact_style == 'label':
              torch.save(args.pseudo_labels, 'pseudo_labels.pth')
              break
          
          # only dim=1000 evaluate unsupervised classification
          if args.dim == 1000:
              eval_stats = eval_one_epoch(args, model, data_loader_val, device, 
                                          logfn=os.path.join(output_dir, 'detail_log.txt'))
              print(eval_stats)
          
          # saving checkpoint
          checkpoint_path = os.path.join(output_dir, 'checkpoint.pth')
          save_dict = {
                  'model': model_without_ddp.state_dict(),
                  'backbone': model_without_ddp.backbone_weights(),
                  'optimizer': optimizer.state_dict(),
                  'epoch': epoch,
                  'args': args,
          }
          if args.use_momentum_encoder:
              save_dict.update({
                  'teacher': teacher_model_without_ddp.state_dict(),
                  'teacher_backbone': teacher_model_without_ddp.backbone_weights(),
              })
          if loss_scaler: # save state dict of loss scaler if using mix precision.
              save_dict.update({'scaler': loss_scaler.state_dict()})
          utils.save_on_master(save_dict, checkpoint_path)
  
          if (epoch + 1) == args.mme_epochs:
              utils.save_on_master(save_dict, os.path.join(output_dir,'mme_ckpt.pth'))
          
          # logging train stats
          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:
                  if epoch == 0:  # write the arguments parameters
                      f.write(json.dumps(args.__dict__) + "\n")
                  f.write(json.dumps(log_stats) + "\n")
                  for k, v in train_stats.items():
                      writer.add_scalar(k, v, epoch)
  
      total_time = time.time() - start_time
      print('Training time {}'.format(str(datetime.timedelta(seconds=int(total_time)))))
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser('TWIST', parents=[get_args_parser()])
      args = parser.parse_args()
      if args.output_dir:
          Path(args.output_dir).mkdir(parents=True, exist_ok=True)
      main(args)
  