text: from __future__ import print_function
  
  import os
  import sys
  import argparse
  import time
  import math
  
  import numpy
  import torch
  import torch.backends.cudnn as cudnn
  
  from main_ce import set_loader
  from util import AverageMeter
  from util import adjust_learning_rate, warmup_learning_rate, accuracy
  from util import set_optimizer
  from util import save_h5
  from tqdm import tqdm
  from networks.resnet_big import SupConResNet, LinearClassifier
  
  try:
      import apex
      from apex import amp, optimizers
  except ImportError:
      pass
  
  
  def parse_option():
      parser = argparse.ArgumentParser('argument for training')
  
      parser.add_argument('--get_embeddings', action='store_true')
      parser.add_argument('--print_freq', type=int, default=10,
                          help='print frequency')
      parser.add_argument('--save_freq', type=int, default=50,
                          help='save frequency')
      parser.add_argument('--batch_size', type=int, default=256,
                          help='batch_size')
      parser.add_argument('--num_workers', type=int, default=16,
                          help='num of workers to use')
      parser.add_argument('--epochs', type=int, default=100,
                          help='number of training epochs')
  
      # optimization
      parser.add_argument('--learning_rate', type=float, default=0.1,
                          help='learning rate')
      parser.add_argument('--lr_decay_epochs', type=str, default='60,75,90',
                          help='where to decay lr, can be a list')
      parser.add_argument('--lr_decay_rate', type=float, default=0.2,
                          help='decay rate for learning rate')
      parser.add_argument('--weight_decay', type=float, default=0,
                          help='weight decay')
      parser.add_argument('--momentum', type=float, default=0.9,
                          help='momentum')
      parser.add_argument('-gpu', '--gpu_ids', default='', help="gpu ids used to train")  # before: default="0,1,2,3"
      parser.add_argument('--val_type', type=str, default='val1_small', help='val[1234][_small] for hotels50k')
  
  
      # model dataset
      parser.add_argument('--model', type=str, default='resnet50')
      parser.add_argument('--dataset', type=str, default='cifar10',
                          choices=['cifar10', 'cifar100', 'hotels'], help='dataset')
      parser.add_argument('--data_folder', type=str, default=None, help='path to custom dataset')
  
      # other setting
      parser.add_argument('--cosine', action='store_true',
                          help='using cosine annealing')
      parser.add_argument('--warm', action='store_true',
                          help='warm-up for large batch training')
  
      parser.add_argument('--ckpt', type=str, default='',
                          help='path to pre-trained model')
  
      opt = parser.parse_args()
  
      # set the path according to the environment
      if opt.data_folder is None:
          opt.data_folder = './datasets/'
  
      iterations = opt.lr_decay_epochs.split(',')
      opt.lr_decay_epochs = list([])
      for it in iterations:
          opt.lr_decay_epochs.append(int(it))
  
      opt.model_name = '{}_{}_lr_{}_decay_{}_bsz_{}'.\
          format(opt.dataset, opt.model, opt.learning_rate, opt.weight_decay,
                 opt.batch_size)
  
      if opt.cosine:
          opt.model_name = '{}_cosine'.format(opt.model_name)
  
      # warm-up for large-batch training,
      if opt.warm:
          opt.model_name = '{}_warm'.format(opt.model_name)
          opt.warmup_from = 0.01
          opt.warm_epochs = 10
          if opt.cosine:
              eta_min = opt.learning_rate * (opt.lr_decay_rate ** 3)
              opt.warmup_to = eta_min + (opt.learning_rate - eta_min) * (
                      1 + math.cos(math.pi * opt.warm_epochs / opt.epochs)) / 2
          else:
              opt.warmup_to = opt.learning_rate
  
      if opt.dataset == 'cifar10':
          opt.n_cls = 10
      elif opt.dataset == 'cifar100':
          opt.n_cls = 100
      elif opt.dataset == 'hotels':
          opt.n_cls = 2468
      else:
          raise ValueError('dataset not supported: {}'.format(opt.dataset))
  
      return opt
  
  
  def set_model(opt):
      model = SupConResNet(name=opt.model)
      criterion = torch.nn.CrossEntropyLoss()
  
      classifier = LinearClassifier(name=opt.model, num_classes=opt.n_cls)
  
      ckpt = torch.load(opt.ckpt, map_location='cpu')
      state_dict = ckpt['model']
  
      if torch.cuda.is_available():
          if torch.cuda.device_count() > 1:
              model.encoder = torch.nn.DataParallel(model.encoder)
          else:
              new_state_dict = {}
              for k, v in state_dict.items():
                  k = k.replace("module.", "")
                  new_state_dict[k] = v
              state_dict = new_state_dict
          model = model.cuda()
          classifier = classifier.cuda()
          criterion = criterion.cuda()
          cudnn.benchmark = True
  
          model.load_state_dict(state_dict)
  
      return model, classifier, criterion
  
  
  def train(train_loader, model, classifier, criterion, optimizer, epoch, opt):
      """one epoch training"""
      model.eval()
      classifier.train()
  
      batch_time = AverageMeter()
      data_time = AverageMeter()
      losses = AverageMeter()
      top1 = AverageMeter()
      top5 = AverageMeter()
  
      end = time.time()
      for idx, (images, labels) in enumerate(train_loader): # todo
          data_time.update(time.time() - end)
  
          images = images.cuda(non_blocking=True)
          labels = labels.cuda(non_blocking=True)
          bsz = labels.shape[0]
  
          # warm-up learning rate
          warmup_learning_rate(opt, epoch, idx, len(train_loader), optimizer)
  
          # compute loss
          with torch.no_grad():
              features = model.encoder(images)
          output = classifier(features.detach())
          loss = criterion(output, labels)
  
          # update metric
          losses.update(loss.item(), bsz)
          acc1, acc5 = accuracy(output, labels, topk=(1, 5))
          top1.update(acc1[0], bsz)
          top5.update(acc5[0], bsz)
  
          # SGD
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          # measure elapsed time
          batch_time.update(time.time() - end)
          end = time.time()
  
          # print info
          if (idx + 1) % opt.print_freq == 0:
              print('Train: [{0}][{1}/{2}]\t'
                    'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                    'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
                    'loss {loss.val:.3f} ({loss.avg:.3f})\t'
                    'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                    'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                     epoch, idx + 1, len(train_loader), batch_time=batch_time,
                     data_time=data_time, loss=losses, top1=top1, top5=top5))
              sys.stdout.flush()
  
      return losses.avg, top1.avg
  
  def get_embs(loader, model, loader_type=''):
      """getting all embeddings"""
      model.eval()
  
      total_num_imgs = loader.dataset.__len__()
  
      all_embs = numpy.zeros((total_num_imgs, 2048), dtype=numpy.float32)  # only for resnet50 and resnet101
      all_lbls = numpy.zeros((total_num_imgs,), dtype=numpy.float32)
  
      with torch.no_grad():
          with tqdm(total=len(loader), desc=f'Saving embs for {loader_type}...') as t:
              for idx, (images, labels) in enumerate(loader):
                  images = images.float().cuda()
                  labels = labels.cuda()
                  bsz = labels.shape[0]
  
                  # forward
                  embs = model.encoder(images)
  
                  begin_idx = idx * bsz
                  end_idx = min((idx + 1) * bsz, total_num_imgs)
  
                  all_embs[begin_idx: end_idx, :] = embs.cpu().detach().numpy()
                  all_lbls[begin_idx: end_idx] = labels.cpu().numpy()
                  t.update()
  
      return all_embs, all_lbls
  
  
  def validate(val_loader, model, classifier, criterion, opt):
      """validation"""
      model.eval()
      classifier.eval()
  
      batch_time = AverageMeter()
      losses = AverageMeter()
      top1 = AverageMeter()
  
      with torch.no_grad():
          end = time.time()
          for idx, (images, labels) in enumerate(val_loader):
              images = images.float().cuda()
              labels = labels.cuda()
              bsz = labels.shape[0]
  
              # forward
              embs = model.encoder(images)
              output = classifier(embs)
              loss = criterion(output, labels)
  
              # update metric
              losses.update(loss.item(), bsz)
              acc1, acc5 = accuracy(output, labels, topk=(1, 5))
              top1.update(acc1[0], bsz)
  
              # measure elapsed time
              batch_time.update(time.time() - end)
              end = time.time()
  
              if idx % opt.print_freq == 0:
                  print('Test: [{0}/{1}]\t'
                        'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                        'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                        'Acc@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                         idx, len(val_loader), batch_time=batch_time,
                         loss=losses, top1=top1))
  
      print(' * Acc@1 {top1.avg:.3f}'.format(top1=top1))
      return losses.avg, top1.avg
  
  
  def main():
      best_acc = 0
      opt = parse_option()
      if opt.gpu_ids != '':
          os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu_ids
      # build data loader
      train_loader, val_loader = set_loader(opt)
  
      # build model and criterion
      model, classifier, criterion = set_model(opt)
  
      # build optimizer
      optimizer = set_optimizer(opt, classifier)
  
      if opt.get_embeddings:
          embs, lbls = get_embs(val_loader, model, opt.val_type)
          save_path = os.path.split(opt.ckpt)[0]
          save_h5('data', lbls, 'i8',
                  os.path.join(save_path, f'{opt.dataset}_{opt.val_type}_Classes.h5'))
          save_h5('data', embs, 'f',
                  os.path.join(save_path, f'{opt.dataset}_{opt.val_type}_Feats.h5'))
      else:
          # training routine
          for epoch in range(1, opt.epochs + 1):
              adjust_learning_rate(opt, optimizer, epoch)
  
              # train for one epoch
              time1 = time.time()
              loss, acc = train(train_loader, model, classifier, criterion,
                                optimizer, epoch, opt)
              time2 = time.time()
              print('Train epoch {}, total time {:.2f}, accuracy:{:.2f}'.format(
                  epoch, time2 - time1, acc))
  
              # eval for one epoch
              loss, val_acc = validate(val_loader, model, classifier, criterion, opt)
              if val_acc > best_acc:
                  best_acc = val_acc
  
          print('best accuracy: {:.2f}'.format(best_acc))
  
  
  if __name__ == '__main__':
      main()
  