text: import argparse
  from cmath import isnan
  import datetime
  import logging
  import os
  import random
  import sys
  
  sys.path.append(".")
  
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.distributed as dist
  import torch.nn.functional as F
  from datasets import voc2 as voc
  from model.losses import get_masked_cos_loss, get_seg_loss, DINOLoss
  from model.model_aux_dino import network
  from omegaconf import OmegaConf
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import DataLoader
  from torch.utils.data.distributed import DistributedSampler
  from torch.utils.tensorboard import SummaryWriter
  from tqdm import tqdm
  from utils import evaluate, imutils, optimizer
  from utils.camutils import cam_to_label, cam_to_label2, multi_scale_cam2, label_to_aff_mask
  from utils.pyutils import AverageMeter, cal_eta, format_tabs, setup_logger, get_mask_by_radius
  
  parser = argparse.ArgumentParser()
  parser.add_argument("--config", default='configs/voc.yaml', type=str, help="config")
  parser.add_argument("--work_dir", default=None, type=str, help="work_dir")
  parser.add_argument("--local_rank", default=-1, type=int, help="local_rank")
  parser.add_argument("--crop_size", default=224, type=int, help="crop_size")
  parser.add_argument('--backend', default='nccl')
  
  def worker_init_fn(worker_id):
      worker_seed = torch.initial_seed() % 2**32
      np.random.seed(worker_seed)
      random.seed(worker_seed)
  
  def setup_seed(seed):
      torch.manual_seed(seed)
      torch.cuda.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
      np.random.seed(seed)
      random.seed(seed)
  
      torch.backends.cudnn.benchmark = False
      torch.backends.cudnn.deterministic = True
  
  
  def validate(model=None, data_loader=None, cfg=None):
  
      preds, gts, cams, tscams, cams_aux = [], [], [], [], []
      model.eval()
      avg_meter = AverageMeter()
      with torch.no_grad():
          for _, data in tqdm(enumerate(data_loader), total=len(data_loader), ncols=100, ascii=" >="):
  
              name, inputs, labels, cls_label = data
              inputs = inputs.cuda()
  
              b, c, h, w = inputs.shape
              inputs  = F.interpolate(inputs, size=[cfg.dataset.crop_size, cfg.dataset.crop_size], mode='bilinear', align_corners=False)
              labels = labels.cuda()
              cls_label = cls_label.cuda()
  
              cls, segs, _, _, _ = model(inputs,)
  
              cls_pred = (cls>0).type(torch.int16)
              _f1 = evaluate.multilabel_score(cls_label.cpu().numpy()[0], cls_pred.cpu().numpy()[0])
              avg_meter.add({"cls_score": _f1})
  
              resized_segs = F.interpolate(segs, size=labels.shape[1:], mode='bilinear', align_corners=False)
              ###
              _cams, _tscam, _cams_aux = multi_scale_cam2(model, inputs, cfg.cam.scales)
              resized_cam = F.interpolate(_cams, size=labels.shape[1:], mode='bilinear', align_corners=False)
              cam_label = cam_to_label(resized_cam, cls_label, cfg=cfg)
  
              resized_cam_aux = F.interpolate(_cams_aux, size=labels.shape[1:], mode='bilinear', align_corners=False)
              cam_label_aux = cam_to_label(resized_cam_aux, cls_label, cfg=cfg)
  
              resized_tscam = F.interpolate(_tscam, size=labels.shape[1:], mode='bilinear', align_corners=False)
              tscam_label = cam_to_label2(resized_tscam, cls_label, bkg_score=0.15, delta=0.1)
  
              cls_pred = (cls > 0).type(torch.int16)
              _f1 = evaluate.multilabel_score(cls_label.cpu().numpy()[0], cls_pred.cpu().numpy()[0])
              avg_meter.add({"cls_score": _f1})
  
              resized_segs = F.interpolate(segs, size=labels.shape[1:], mode='bilinear', align_corners=False)
  
              preds += list(torch.argmax(resized_segs, dim=1).cpu().numpy().astype(np.int16))
              cams += list(cam_label.cpu().numpy().astype(np.int16))
              gts += list(labels.cpu().numpy().astype(np.int16))
              cams_aux += list(cam_label_aux.cpu().numpy().astype(np.int16))
              tscams += list(tscam_label.cpu().numpy().astype(np.int16))
  
              valid_label = torch.nonzero(cls_label[0])[:, 0]
              out_cam = torch.squeeze(resized_cam)[valid_label]
              #np.save(os.path.join(cfg.work_dir.pred_dir, name[0]+'.npy'), {"keys":valid_label.cpu().numpy(), "cam":out_cam.cpu().numpy()})
  
      cls_score = avg_meter.pop('cls_score')
      seg_score = evaluate.scores(gts, preds)
      cam_score = evaluate.scores(gts, cams)
      cam_aux_score = evaluate.scores(gts, cams_aux)
      tscam_score = evaluate.scores(gts, tscams)
      model.train()
      tab_results = format_tabs([cam_score, cam_aux_score, tscam_score, seg_score], name_list=["CAM", "aux_CAM", "TS-CAM", "Segs"], cat_list=voc.class_list)
  
      return cls_score, tab_results
  
  def train(cfg):
  
      num_workers = 10
  
      torch.cuda.set_device(args.local_rank)
      dist.init_process_group(backend=args.backend, )
  
      time0 = datetime.datetime.now()
      time0 = time0.replace(microsecond=0)
  
      train_dataset = voc.VOC12ClsDataset(
          root_dir=cfg.dataset.root_dir,
          name_list_dir=cfg.dataset.name_list_dir,
          split=cfg.train.split,
          stage='train',
          aug=True,
          resize_range=cfg.dataset.resize_range,
          rescale_range=cfg.dataset.rescale_range,
          crop_size=cfg.dataset.crop_size,
          img_fliplr=True,
          ignore_index=cfg.dataset.ignore_index,
          num_classes=cfg.dataset.num_classes,
      )
  
      val_dataset = voc.VOC12SegDataset(
          root_dir=cfg.dataset.root_dir,
          name_list_dir=cfg.dataset.name_list_dir,
          split=cfg.val.split,
          stage='val',
          aug=False,
          ignore_index=cfg.dataset.ignore_index,
          num_classes=cfg.dataset.num_classes,
      )
  
      train_sampler = DistributedSampler(train_dataset, shuffle=True)
      train_loader = DataLoader(
          train_dataset,
          batch_size=cfg.train.samples_per_gpu,
          #shuffle=True,
          num_workers=num_workers,
          pin_memory=False,
          drop_last=True,
          sampler=train_sampler,
          prefetch_factor=4)
  
      val_loader = DataLoader(val_dataset,
                              batch_size=1,
                              shuffle=False,
                              num_workers=num_workers,
                              pin_memory=False,
                              drop_last=False)
  
      device = torch.device(args.local_rank)
  
      model = network(
          backbone=cfg.backbone.name,
          num_classes=cfg.dataset.num_classes,
          embedding_dim=256,
          pretrained=cfg.backbone.pretrained,
          pooling=cfg.backbone.pooling,
          total_iters=cfg.train.cam_iters,
          init_momentum=cfg.train.dino_momentum
      )
      param_groups = model.get_param_groups()
      model.to(device)
  
      if args.local_rank == 0:
          writer = SummaryWriter(cfg.work_dir.tb_logger_dir)
  
      optim = getattr(optimizer, cfg.optimizer.type)(
          params=[
              {
                  "params": param_groups[0],
                  "lr": cfg.optimizer.learning_rate,
                  "weight_decay": cfg.optimizer.weight_decay,
              },
              {
                  "params": param_groups[1],
                  "lr": cfg.optimizer.learning_rate,
                  "weight_decay": cfg.optimizer.weight_decay,
              },
              {
                  "params": param_groups[2],
                  "lr": cfg.optimizer.learning_rate * 10,
                  "weight_decay": cfg.optimizer.weight_decay,
              },
              {
                  "params": param_groups[3],
                  "lr": cfg.optimizer.learning_rate * 10,
                  "weight_decay": cfg.optimizer.weight_decay,
              },
          ],
          lr=cfg.optimizer.learning_rate,
          weight_decay=cfg.optimizer.weight_decay,
          betas=cfg.optimizer.betas,
          warmup_iter=cfg.scheduler.warmup_iter,
          max_iter=cfg.train.max_iters,
          warmup_ratio=cfg.scheduler.warmup_ratio,
          power=cfg.scheduler.power)
  
      logging.info('\nOptimizer: \n%s' % optim)
      model = DistributedDataParallel(model, device_ids=[args.local_rank], find_unused_parameters=True)
  
      train_sampler.set_epoch(np.random.randint(cfg.train.max_iters))
      train_loader_iter = iter(train_loader)
      # loss_mask = get_mask_by_radius(h=cfg.dataset.crop_size, w=cfg.dataset.crop_size, radius=4)
  
      avg_meter = AverageMeter()
  
      # fake_iter_per_epoch = 800
      # total_epochs = int(cfg.train.max_iters / fake_iter_per_epoch)
      DINO_loss = DINOLoss(
          out_dim=1024,
          ncrops=10,).cuda()
  
      for n_iter in range(cfg.train.max_iters):
  
          try:
              img_name, inputs, cls_label, img_box, crops = next(train_loader_iter)
  
          except:
              train_sampler.set_epoch(np.random.randint(cfg.train.max_iters))
              train_loader_iter = iter(train_loader)
              img_name, inputs, cls_label, img_box, crops = next(train_loader_iter)
  
          inputs = inputs.to(device, non_blocking=True)
          cls_label = cls_label.to(device, non_blocking=True)
  
          cls, segs, attns, fmap, cls_aux, out_t, out_s = model(inputs, crops=crops, n_iter=n_iter)
  
          cams, tscams, cams_aux = multi_scale_cam2(model, inputs=inputs, scales=cfg.cam.scales)
          valid_cam, pseudo_label = cam_to_label(cams.detach(), cls_label=cls_label, img_box=img_box, ignore_mid=True, cfg=cfg)
          valid_cam_aux, _ = cam_to_label(cams_aux.detach(), cls_label=cls_label, img_box=img_box, ignore_mid=True, cfg=cfg)
  
          segs = F.interpolate(segs, size=pseudo_label.shape[1:], mode='bilinear', align_corners=False)
  
          seg_loss = get_seg_loss(segs, pseudo_label.type(torch.long), ignore_index=cfg.dataset.ignore_index)
          cls_loss = F.multilabel_soft_margin_loss(cls, cls_label)
          cls_loss_aux = F.multilabel_soft_margin_loss(cls_aux, cls_label)
          di_loss = DINO_loss(out_s, out_t,)
  
          resized_cams_aux = F.interpolate(cams_aux, size=fmap.shape[2:], mode="bilinear", align_corners=False)
          _, pseudo_label_aux = cam_to_label(resized_cams_aux.detach(), cls_label=cls_label, img_box=img_box, ignore_mid=True, cfg=cfg)
          aff_mask = label_to_aff_mask(pseudo_label_aux)
  
          cos_loss = get_masked_cos_loss(fmap, aff_mask)
          # cos_loss = get_cos_loss(fmap, low_fmap)
  
          loss = 1.0 * cls_loss + 0.0 * seg_loss + 1.0 * cls_loss_aux + 0.5 * cos_loss + 0.5 * di_loss
  
          cls_pred = (cls > 0).type(torch.int16)
          cls_score = evaluate.multilabel_score(cls_label.cpu().numpy()[0], cls_pred.cpu().numpy()[0])
  
          avg_meter.add({'cls_loss': cls_loss.item(), 'cos_loss': cos_loss.item(), 'dino_loss': di_loss.item(), 'cls_loss_aux': cls_loss_aux.item(), 'cls_score': cls_score.item(),})
  
          optim.zero_grad()
          loss.backward()
          optim.step()
  
          if (n_iter + 1) % cfg.train.log_iters == 0:
  
              delta, eta = cal_eta(time0, n_iter + 1, cfg.train.max_iters)
              cur_lr = optim.param_groups[0]['lr']
  
              preds = torch.argmax(segs, dim=1,).cpu().numpy().astype(np.int16)
              gts = pseudo_label.cpu().numpy().astype(np.int16)
  
              seg_mAcc = (preds == gts).sum() / preds.size
  
              grid_imgs, grid_cam = imutils.tensorboard_image(imgs=inputs.clone(), cam=valid_cam)
              _, grid_tscam = imutils.tensorboard_image(imgs=inputs.clone(), cam=tscams)
              _, grid_cam_aux = imutils.tensorboard_image(imgs=inputs.clone(), cam=valid_cam_aux)
  
              _attns_detach = [a.detach() for a in attns]
  
              # grid_attns = imutils.tensorboard_attn2(attns=_attns_detach, n_row=cfg.train.samples_per_gpu, with_attn_pred=False)
              grid_attn_cls_mhsa, grid_attn_cls_mean = imutils.tensorboard_cls_token_attn(_attns_detach)
  
              grid_labels = imutils.tensorboard_label(labels=gts)
              grid_preds = imutils.tensorboard_label(labels=preds)
  
              if args.local_rank == 0:
                  logging.info("Iter: %d; Elasped: %s; ETA: %s; LR: %.3e; cls_loss: %.4f, cls_loss_aux: %.4f, cos_loss %.4f, dino_loss %.4f" % (n_iter + 1, delta, eta, cur_lr, avg_meter.pop('cls_loss'), avg_meter.pop('cls_loss_aux'), avg_meter.pop('cos_loss'), avg_meter.pop('dino_loss')))
  
                  writer.add_image("train/images", grid_imgs, global_step=n_iter)
                  writer.add_image("train/preds", grid_preds, global_step=n_iter)
                  writer.add_image("train/pseudo_gts", grid_labels, global_step=n_iter)
  
                  writer.add_image("cam/valid_cams", grid_cam, global_step=n_iter)
                  writer.add_image("cam/valid_cams_aux", grid_cam_aux, global_step=n_iter)
                  writer.add_image("cam/valid_tscams", grid_tscam, global_step=n_iter)
  
                  writer.add_image("attns/cls_token_mean_per_layer", grid_attn_cls_mean, global_step=n_iter)
                  writer.add_image("attns/cls_token_multi-head_last_layer", grid_attn_cls_mhsa, global_step=n_iter)
  
                  writer.add_scalars('loss', {"seg_loss": seg_loss.item(), "cls_loss": cls_loss.item()}, global_step=n_iter)
                  writer.add_scalars('learning_rate', {"learning_rate": optim.param_groups[0]["lr"], }, global_step=n_iter)
  
          if (n_iter + 1) % cfg.train.eval_iters == 0:
              ckpt_name = os.path.join(cfg.work_dir.ckpt_dir, "model_iter_%d.pth" % (n_iter + 1))
              if args.local_rank == 0:
                  logging.info('Validating...')
                  torch.save(model.state_dict(), ckpt_name)
              val_cls_score, tab_results = validate(model=model, data_loader=val_loader, cfg=cfg)
              if args.local_rank == 0:
                  logging.info("val cls score: %.6f" % (val_cls_score))
                  logging.info("\n"+tab_results)
  
                  writer.add_scalars('cls_score', {"val": val_cls_score, "train": avg_meter.pop('cls_score')}, global_step=n_iter)
  
      return True
  
  
  if __name__ == "__main__":
  
      args = parser.parse_args()
      cfg = OmegaConf.load(args.config)
  
      if args.work_dir is not None:
          cfg.work_dir.dir = args.work_dir
          
      if args.crop_size is not None:
          cfg.dataset.crop_size = args.crop_size
  
      timestamp = "{0:%Y-%m-%d-%H-%M-%S}".format(datetime.datetime.now())
      cfg.work_dir.dir = os.path.join(cfg.work_dir.dir, timestamp)
      cfg.work_dir.ckpt_dir = os.path.join(cfg.work_dir.dir, cfg.work_dir.ckpt_dir)
      cfg.work_dir.pred_dir = os.path.join(cfg.work_dir.dir, cfg.work_dir.pred_dir)
      cfg.work_dir.tb_logger_dir = os.path.join(cfg.work_dir.dir, cfg.work_dir.tb_logger_dir)
  
      os.makedirs(cfg.work_dir.ckpt_dir, exist_ok=True)
      os.makedirs(cfg.work_dir.pred_dir, exist_ok=True)
      os.makedirs(cfg.work_dir.tb_logger_dir, exist_ok=True)
  
      if args.local_rank == 0:
          setup_logger(filename=os.path.join(cfg.work_dir.dir, 'train.log'))
          logging.info('Pytorch version: %s' % torch.__version__)
          logging.info("GPU type: %s"%(torch.cuda.get_device_name(0)))
          logging.info('\nargs: %s' % args)
  
          with open(os.path.join(cfg.work_dir.dir, 'config.yaml'), 'w') as f:
              OmegaConf.save(cfg, f.name)
  
      ## fix random seed
      setup_seed(1)
      train(cfg=cfg)
  