text: # from __future__ import print_function
  
  import os,argparse
  import random
  import gc
  import numpy as np
  import cv2
  from PIL import Image
  import time
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  from torchvision import datasets, transforms
  from torch.utils.data.dataset import Dataset
  import torchvision.models as models
  
  from sklearn.model_selection import KFold
  
  
  from libs.mAP import getValmAP
  from libs.tools import *
  from libs.model import NetClassify, NetMultilabel
  from libs.data import getDataLoader
  from libs.mixup import mixup_data, mixup_criterion
  
  from torch.autograd import Variable
  from config import cfg
  
  import platform
  
  from libs.scheduler import GradualWarmupScheduler
  
  import glob
  
  #from adabelief_pytorch import AdaBelief
  
  from libs.ranger import Ranger 
  from libs.focal_loss import FocalLoss 
  
  
  def trainClassify(model, 
      device, 
      train_loader, 
      optimizer, 
      epoch, 
      total_epoch,
      criterion,
      use_distill,
      label_smooth):
      
      model.train()
      correct = 0
      count = 0
      for batch_idx, (data, target) in enumerate(train_loader):
          data, target = data.to(device), target.to(device)
  
          output = model(data).double()
  
          #all_linear2_params = torch.cat([x.view(-1) for x in model.model_feature._fc.parameters()])
          #l2_regularization = 0.0003 * torch.norm(all_linear2_params, 2)
  
          loss = criterion(output, target)# + l2_regularization.item()
          loss.backward() #计算梯度
  
          clip_gradient(optimizer)
  
          optimizer.step() #更新参数
          optimizer.zero_grad()#把梯度置零
  
          ### train acc
          pred_score = nn.Softmax(dim=1)(output)
          pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
          if use_distill or label_smooth>0:
              target = target.max(1, keepdim=True)[1] 
          correct += pred.eq(target.view_as(pred)).sum().item()
          count += len(data)
  
          train_acc =  correct / count
          #print(train_acc)
          if batch_idx % 10 == 0:
              print('\r',
                  '{}/{} [{}/{} ({:.0f}%)] loss:{:.3f} acc: {:.3f} '.format(
                  epoch+1, total_epoch,batch_idx * len(data), len(train_loader.dataset),
                  100. * batch_idx / len(train_loader), loss.item(),train_acc), 
                  end="",flush=True)
  
  
  #val_loss: 1.0412, val_acc: 66.67%, val_mAP: 0.6167
  def valClassify( model, device, val_loader, criterion, use_distill, label_smooth):
      model.eval()
      val_loss = 0
      correct = 0
  
      with torch.no_grad():
          pres = []
          labels = []
          for data, target in val_loader:
              data, target = data.to(device), target.to(device)
              #print(target.shape)
              if use_distill:
                  output = model(data).double()
              else:
                  output = model(data)
  
  
              val_loss += criterion(output, target).item() # sum up batch loss
  
              #print(output.shape)
              pred_score = nn.Softmax(dim=1)(output)
              #print(pred_score.shape)
              pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
              if use_distill or label_smooth>0:
                  target = target.max(1, keepdim=True)[1] 
              correct += pred.eq(target.view_as(pred)).sum().item()
  
  
              batch_pred_score = pred_score.data.cpu().numpy().tolist()
              batch_label_score = target.data.cpu().numpy().tolist()
              pres.extend(batch_pred_score)
              labels.extend(batch_label_score)
  
      pres = np.array(pres)
      labels = np.array(labels)
      #print(pres.shape, labels.shape)
  
  
      mAP = getValmAP(pres, labels)
  
      val_loss /= len(val_loader.dataset)
      val_acc =  correct / len(val_loader.dataset)
      print(' ------------------------------ val_loss: {:.4f}, val_acc: {:.2f}%, val_mAP: {:.4f}'.format(
          val_loss, 100. * val_acc, mAP))
  
      return val_loss,  mAP
  
  
  
  
  
  def main(cfg):
      
  
      print(cfg)
      print("=================================")
  
  
      model_name = cfg['model_name']
      img_size = cfg['img_size']
      class_number = cfg['class_number']
      save_dir = cfg['save_dir']
      random_seed = cfg['random_seed']
      train_path = cfg['train_path']
      GPU_ID = cfg['GPU_ID']
  
      fold_num = cfg['k_flod']
      batch_size = cfg['batch_size']
      epochs = cfg['epochs']
      learning_rate = cfg['learning_rate']
      early_stop_patient = cfg['early_stop_patient']
      save_start_epoch = cfg['save_start_epoch']
      use_warmup = cfg['use_warmup']
      schedu = cfg['schedu']
      optims = cfg['optims']
      weight_decay = cfg['weight_decay']
      use_distill = cfg['use_distill']
      label_smooth = cfg['label_smooth']
      model_path = cfg['model_path']
      start_fold = cfg['start_fold']
  
      os.environ["CUDA_VISIBLE_DEVICES"] = GPU_ID
      seed_reproducer(random_seed)
  
      ############################################################
      
  
      # log_interval = 10
      #use_cuda = True
      device = torch.device("cuda")#cuda
  
      if platform.system() == "Windows":
          kwargs = {'num_workers': 0, 'pin_memory': True}
      else:
          kwargs = {'num_workers': 4, 'pin_memory': True}
  
      train_names = getAllName(train_path)
      print("total imgs: ", len(train_names))
  
      if not use_distill:
          train_names = [x for x in train_names if "aug" not in x]
          print("remove aug: ", len(train_names))
  
      #print(train_names[:3])
      train_names.sort(key = lambda x:os.path.basename(x))
      #print(train_names[:3])
  
      train_names = np.array(train_names)
      random.shuffle(train_names)
  
   
  
      
      folds = KFold(n_splits=fold_num, shuffle=False)#, random_state=random_seed
      for fold_i, (train_index, val_index) in enumerate(folds.split(train_names)):
          print("Fold: ", fold_i+1,'/',fold_num)
          if fold_i<start_fold:
              continue
  
  
          train_data = train_names[train_index]
          val_data = train_names[val_index]
          # print(val_data[-3:])
          # b
          input_data = [train_data, val_data]
          
  
  
          if not use_distill and label_smooth==0:
              criterion = torch.nn.CrossEntropyLoss().cuda()
              #criterion = FocalLoss().cuda()
              train_loader, val_loader = getDataLoader("trainClassify", input_data,model_name, img_size, batch_size, kwargs)
          else:
  
              criterion = CrossEntropyLossOneHot().cuda()
              #kwargs['use_distill'] = use_distill
              #print(kwargs)
              train_loader, val_loader = getDataLoader("trainClassifyOnehot", input_data,model_name, img_size, batch_size, kwargs)
  
  
          model = NetClassify(model_name, class_number).to(device)
          if model_path is not None:
              model.load_state_dict(torch.load(model_path))
              print("---------------------- load model!!!")
          # print(model)
          # b
          
  
          if optims=='adam':
              optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
          elif optims=='SGD':
              optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=weight_decay)
          elif optims=='AdaBelief':
              optimizer = AdaBelief(model.parameters(), lr=learning_rate, eps=1e-12, betas=(0.9,0.999))
          elif optims=='Ranger':
              optimizer = Ranger(model.parameters(), lr=learning_rate)
          
  
  
          if schedu=='default':
              scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.1, patience=5)
          elif schedu=='step1':
              scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.8, last_epoch=-1)
          elif schedu=='step2':
              scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=8, gamma=0.5, last_epoch=-1)
          elif schedu=='step3':
              scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5, last_epoch=-1)
          elif schedu=='SGDR1': 
              scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer,
                                                                   T_0=10, 
                                                                  T_mult=2)
          elif schedu=='SGDR2': 
              scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer,
                                                                   T_0=5, 
                                                                  T_mult=2)
  
          # elif schedu=='CVPR': 
          #     scheduler = WarmRe   mizer, T_max=10, T_mult=1, eta_min=1e-5)
          
          if use_warmup:
              scheduler_warmup = GradualWarmupScheduler(optimizer, 
                  multiplier=1, total_epoch=1, after_scheduler=scheduler)
  
  
  
  
  
          early_stop_value = 0
          early_stop_dist = 0
  
          for epoch in range(epochs):
              
              if schedu=='step3':
                  if epoch==10:
                      img_size=416
                      batch_size = 4
                      scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.5, last_epoch=-1)
                      if not use_distill:
                          criterion = torch.nn.CrossEntropyLoss().cuda()
                          train_loader, val_loader = getDataLoader("trainClassify", input_data,model_name, img_size, batch_size, kwargs)
                      else:
                          criterion = CrossEntropyLossOneHot().cuda()
                          train_loader, val_loader = getDataLoader("trainClassifyOnehot", input_data,model_name, img_size, batch_size, kwargs)
  
  
                  elif epoch==15:
                      img_size=600
                      batch_size = 3
                      scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer,
                                                                       T_0=10, 
                                                                      T_mult=2)
                      if not use_distill:
                          criterion = torch.nn.CrossEntropyLoss().cuda()
                          train_loader, val_loader = getDataLoader("trainClassify", input_data,model_name, img_size, batch_size, kwargs)
                      else:
                          criterion = CrossEntropyLossOneHot().cuda()
                          train_loader, val_loader = getDataLoader("trainClassifyOnehot", input_data,model_name, img_size, batch_size, kwargs)
  
  
              
  
              trainClassify(model, device, train_loader, optimizer, epoch, epochs, criterion, use_distill, label_smooth)
              print(" LR:", optimizer.param_groups[0]["lr"], end="")
  
              t = time.time()
              val_loss, mAP = valClassify(model, device, val_loader, criterion, use_distill, label_smooth)
              print("val time: ", time.time() - t)
  
              #print('333')
              #continue
  
              if use_warmup:
                  scheduler_warmup.step(epoch)
              else:
                  if schedu=='default':
                      scheduler.step(mAP)
                  else:
                      scheduler.step()
  
  
              #print("---")
              #print(mAP, early_stop_value, early_stop_dist)
              if mAP>early_stop_value:
                  early_stop_value = mAP
                  early_stop_dist = 0
                  if epoch>=save_start_epoch:
                      hitory_path = glob.glob('./save/%s-%d_*k-%d_%s.pth' % (model_name,img_size,fold_i,GPU_ID))
                      if len(hitory_path)!=0:
                          if os.path.exists(hitory_path[0]):
                              os.remove(hitory_path[0])
                      torch.save(model.state_dict(), './save/%s-%d_%d_%.4f_k-%d_%s.pth' % (model_name,img_size,epoch,mAP,fold_i,GPU_ID))
              
  
  
  
              early_stop_dist+=1
              if early_stop_dist>early_stop_patient:
                  print("------")
                  print(cfg)
                  print("------")
                  print("===== Early Stop with patient %d , best is Epoch - %d :%f" % (early_stop_patient,epoch-early_stop_patient,early_stop_value))
                  break
              if  epoch+1==epochs:
                  print("===== Finish trainging , best is Epoch - %d :%f" % (epoch-early_stop_dist,early_stop_value))
                  break
  
              
  
  
                  
          del model
          gc.collect()
          torch.cuda.empty_cache()
  
  
          #if not use_distill:
          #    break
          #break
  
  
  if __name__ == '__main__':
      main(cfg)