text: import os
  import gc
  import argparse
  import numpy as np
  import logging as logger
  import argparse
  
  import torch
  import torch.nn as nn
  import torch.distributed as dist
  import torch.utils.data.distributed
  import torch.backends.cudnn as cudnn
  from apex.parallel import DistributedDataParallel
  from apex.parallel import convert_syncbn_model
  
  from datasets import make_dataloader
  from config import config as cfg
  from models import make_model
  from losses import SoftLoss, SP_KD_Loss
  from utils import write_config_into_log, ramp_up, ramp_down
  
  
  parser = argparse.ArgumentParser(description="DDP parameters")
  parser.add_argument('--local_rank', type=int, default=0)
  parser.add_argument('--sync_bn', action='store_true', help='enabling apex sync BN')
  
  args = parser.parse_args()
  
  
  cfg.output_dir   = os.path.join(cfg.ckpt_root_dir, cfg.output_dir)
  cfg.snapshot_dir = os.path.join(cfg.output_dir, 'snapshots')
  cfg.train_log    = os.path.join(cfg.output_dir, 'train_log.txt')
  cfg.test_log     = os.path.join(cfg.output_dir, 'test_log.txt')
  
  if args.local_rank == 0:
      print('Experiments dir: {}'.format(cfg.output_dir))
      if not os.path.exists(cfg.output_dir):
          os.makedirs(cfg.output_dir)
      if not os.path.exists(cfg.snapshot_dir):
          os.makedirs(cfg.snapshot_dir)
  
      log_format = '%(levelname)s %(asctime)s %(filename)s] %(message)s'
      logger.basicConfig(level=logger.INFO, format=log_format, datefmt='%Y-%m-%d %H:%M:%S')
      fh = logger.FileHandler(cfg.train_log)
      fh.setFormatter(logger.Formatter(log_format))
      logger.getLogger().addHandler(fh)
  
  
  def train():
      if args.local_rank == 0:
          logger.info('Initializing....')
      cudnn.enable = True
      cudnn.benchmark = True
      # torch.manual_seed(1)
      # torch.cuda.manual_seed(1)
      args.distributed = False
      if 'WORLD_SIZE' in os.environ:
          args.distributed = int(os.environ['WORLD_SIZE']) > 1
      args.gpu = 0
      args.world_size = 1
      if args.distributed:
          args.gpu = args.local_rank
          torch.cuda.set_device(args.gpu)
          torch.distributed.init_process_group(backend='nccl', init_method='env://')
          args.world_size = torch.distributed.get_world_size()
      
      if args.local_rank == 0:
          write_config_into_log(cfg)
  
      if args.local_rank == 0:
          logger.info('Building model......')
      if cfg.pretrained:
          model = make_model(cfg)
          model.load_param(cfg)
          if args.local_rank == 0:
              logger.info('Loaded pretrained model from {0}'.format(cfg.pretrained))
      else:
          model = make_model(cfg)
  
      if args.sync_bn:
          if args.local_rank == 0: logging.info("using apex synced BN")
          model = convert_syncbn_model(model)
      model.cuda()
      if args.distributed:
          # By default, apex.parallel.DistributedDataParallel overlaps communication with
          # computation in the backward pass.
          # delay_allreduce delays all communication to the end of the backward pass.
          model = DistributedDataParallel(model, delay_allreduce=True)
      else:
          model = torch.nn.DataParallel(model)
  
      optimizer = torch.optim.Adam([{'params': model.module.base.parameters(), 'lr': cfg.get_lr(0)[0]},
                                    {'params': model.module.classifiers.parameters(), 'lr': cfg.get_lr(0)[1]}],
                                   weight_decay=cfg.weight_decay)
      
      celoss = nn.CrossEntropyLoss().cuda()
      softloss = SoftLoss()
      sp_kd_loss = SP_KD_Loss()
      criterions = [celoss, softloss, sp_kd_loss]
  
      cfg.batch_size  = cfg.batch_size  // args.world_size
      cfg.num_workers = cfg.num_workers // args.world_size
      train_loader, val_loader = make_dataloader(cfg)
  
      if args.local_rank == 0:
          logger.info('Begin training......')
      for epoch in range(cfg.start_epoch, cfg.max_epoch):
          train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg)
  
          total_acc = test(cfg, val_loader, model, epoch)
          if args.local_rank == 0:
              with open(cfg.test_log, 'a+') as f:
                  f.write('Epoch {0}: Acc is {1:.4f}\n'.format(epoch, total_acc))
              torch.save(obj=model.state_dict(),
                         f=os.path.join(cfg.snapshot_dir, 'ep{}_acc{:.4f}.pth'.format(epoch, total_acc)))
              logger.info('Model saved')
  
  
  def train_one_epoch(train_loader, val_loader, model, criterions, optimizer, epoch, cfg):
      model.train()
      
      ramp_up_w, ramp_down_w = ramp_up(epoch, cfg.ramp_a), ramp_down(epoch, cfg.ramp_a)
      w, gamma = cfg.w, cfg.gamma
      if args.local_rank == 0:
          logger.info('ramp_up_w: {:.4f}, ramp_down_w: {:.4f}, w: {}, gamma: {}'.format(ramp_up_w, ramp_down_w, w, gamma))
      adjust_learning_rate(optimizer, cfg.get_lr(epoch)[0], cfg.get_lr(epoch)[1])
      training_phase = cfg.train_mode
  
      for idx, batch in enumerate(train_loader, start=1):
          img, label = batch
  
          good_batch = check_batch(label, cfg.num_classes)
          good_batch = torch.cuda.IntTensor([good_batch])
          dist.all_reduce(good_batch, op=dist.ReduceOp.SUM)
          good_batch = good_batch // args.world_size
          good_batch = bool(good_batch.item())
          if not good_batch:
              continue
          img, label = img.cuda(), label.cuda()
          
          x_final, output_x_list, targets, softlabel_list, G_matrixs, G_main_matrixs, score, atten_x_final = model(img, label, training_phase)
          
          softlabel = torch.zeros_like(x_final)
          for c in range(cfg.num_classes):
              # sharpen
              px = torch.softmax(softlabel_list[c], dim=1)
              ptx = px ** (1 / cfg.T)  # temparature sharpening
              targets_x = ptx / ptx.sum(dim=1, keepdim=True)  # normalize
              targets_x = targets_x.detach()
  
              _h, _w = px.shape
              _targets = torch.zeros([_h, _w+1]).float()
              _targets[:, 0:c], _targets[:, c + 1:] = targets_x[:, 0:c], targets_x[:, c:]
              ind = (label == c).nonzero()
              softlabel[ind[:, 0]] = _targets.cuda() # bs x num_cls
          
          softLoss = criterions[1](x_final, label, softlabel)
          aux_loss = [criterions[0](_pred, _label) for _pred, _label in zip(output_x_list, targets)] # auxiliary branch  
          aux_loss = sum(aux_loss) / (cfg.num_branches-1)
          CEloss = criterions[0](atten_x_final, label) # main
          spLoss = criterions[2](G_matrixs, G_main_matrixs)
  
          loss = ramp_up_w * CEloss + ramp_up_w * (w * softLoss + gamma * spLoss) +  ramp_down_w * aux_loss
  
          if idx % 20 == 0:
              if args.distributed:
                  CEloss   = reduce_tensor(CEloss)
                  softLoss = reduce_tensor(softLoss)
                  spLoss   = reduce_tensor(spLoss)
                  aux_loss = reduce_tensor(aux_loss)
              if args.local_rank == 0:
                  logger.info('Epoch {} Batch {}/{}: CEloss {:.6f}, softLoss {:.6f}, spLoss {:.6f}, aux_loss {:.6f}'.format(
                                  epoch, idx, len(train_loader), CEloss, softLoss, spLoss, aux_loss))
  
          # compute gradient and do SGD step
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          
          if idx % 200 == 0:
              total_acc = test(cfg, val_loader, model, epoch, idx)
              if args.local_rank == 0:
                  with open(cfg.test_log, 'a+') as f:
                      f.write('Epoch {0} Batch {1}: Acc is {2:.4f}\n'.format(epoch, idx, total_acc))
                  torch.save(obj=model.state_dict(),
                             f=os.path.join(cfg.snapshot_dir, 'ep{}_b{}_acc{:.4f}.pth'.format(epoch, idx, total_acc)))
                  logger.info('Model saved')
  
          gc.collect()
  
  
  def test_worker(val_loader, model):
      Pred, Label = [], []
      for idx, batch in enumerate(val_loader, 1):
          img, label = batch
          pred, output_x_list, targets, softlabel_list = model(img.cuda(), label, 'normal')
  
          for i in range(pred.shape[0]):
              x = pred[i].data.cpu().numpy()
              y = label[i].item()
              Pred.append(x)
              Label.append(y)
  
      return Pred, Label
  
  
  def test(cfg, val_loader, model, epoch, batch=None):
      model.eval()
  
      pred, label = test_worker(val_loader, model)
      pred = [np.where(x == np.max(x))[0][0] for x in pred]
  
      total = len([x for x in range(len(pred)) if pred[x] == label[x]])
      total_acc = total / len(pred)
      if args.distributed:
          total_acc = reduce_float(total_acc)
      if args.local_rank == 0:
          logger.info('Epoch {}: Acc is {:.4f}'.format(epoch, total_acc)) if batch is None else \
              logger.info('Epoch {} Batch {}: Acc is {:.4f}'.format(epoch, batch, total_acc))
  
      return total_acc
  
  
  def adjust_learning_rate(optimizer, lr1, lr2):
      optimizer.param_groups[0]['lr'] = lr1
      optimizer.param_groups[1]['lr'] = lr2
  
  
  def accuracy(output, target, topk=(1,)):
      """Computes the precision@k for the specified values of k"""
      maxk = max(topk)
      batch_size = target.size(0)
  
      _, pred = output.topk(maxk, 1, True, True)
      pred = pred.t()
      correct = pred.eq(target.view(1, -1).expand_as(pred))
  
      res = []
      for k in topk:
          correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
          res.append(correct_k.mul_(100.0 / batch_size))
      return res
  
  
  def check_batch(label, num_classes):
      for i in range(num_classes):
          cnt = (label == i).nonzero().shape[0]
          if cnt < 2:
              return False
      return True
  
  
  def reduce_tensor(tensor):
      rt = tensor.clone()
      dist.all_reduce(rt, op=dist.ReduceOp.SUM)
      rt /= args.world_size
      return rt
  
  
  def reduce_float(f):
      tensor = torch.cuda.FloatTensor([f])
      rt = tensor.clone()
      dist.all_reduce(rt, op=dist.ReduceOp.SUM)
      rt /= args.world_size
      return rt.item()
  
  
  if __name__ == '__main__':
      train()
  