text: import argparse
  import os, sys
  import os.path as osp
  import torchvision
  from torchvision import transforms
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.optim as optim
  import network, loss
  from torch.utils.data import DataLoader
  from data_list import ImageList, ImageList_idx
  import random, pdb, math, copy
  from tqdm import tqdm
  from scipy.spatial.distance import cdist
  from sklearn.metrics import confusion_matrix
  import rotation
  
  def op_copy(optimizer):
      for param_group in optimizer.param_groups:
          param_group['lr0'] = param_group['lr']
      return optimizer
  
  def lr_scheduler(optimizer, iter_num, max_iter, gamma=10, power=0.75):
      decay = (1 + gamma * iter_num / max_iter) ** (-power)
      for param_group in optimizer.param_groups:
          param_group['lr'] = param_group['lr0'] * decay
          param_group['weight_decay'] = 1e-3 
          param_group['momentum'] = 0.9
          param_group['nesterov'] = True
      return optimizer
  
  def image_train(resize_size=256, crop_size=224, alexnet=False):
      normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
      return  transforms.Compose([
          transforms.Resize((resize_size, resize_size)),
          transforms.RandomCrop(crop_size),
          transforms.RandomHorizontalFlip(),
          transforms.ToTensor(),
          normalize
      ])
  
  def image_test(resize_size=256, crop_size=224, alexnet=False):
      normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
      return  transforms.Compose([
          transforms.Resize((resize_size, resize_size)),
          transforms.CenterCrop(crop_size),
          transforms.ToTensor(),
          normalize
      ])
  
  def data_load(args): 
      dsets = {}
      dset_loaders = {}
      train_bs = args.batch_size
      txt_tar = open(args.t_dset_path).readlines()
      txt_test = open(args.test_dset_path).readlines()
  
      dsets["target"] = ImageList_idx(txt_tar, transform=image_train())
      dset_loaders["target"] = DataLoader(dsets["target"], batch_size=train_bs, shuffle=True, num_workers=args.worker, drop_last=False)
      dsets["test"] = ImageList_idx(txt_test, transform=image_test())
      dset_loaders["test"] = DataLoader(dsets["test"], batch_size=train_bs*3, shuffle=False, num_workers=args.worker, drop_last=False)
  
      return dset_loaders
  
  def cal_acc(loader, net, flag=False):
      start_test = True
      with torch.no_grad():
          iter_test = iter(loader)
          for i in range(len(loader)):
              data = iter_test.next()
              inputs = data[0]
              labels = data[1]
              inputs = inputs.cuda()
              _, outputs = net(inputs)
              if start_test:
                  all_output = outputs.float().cpu()
                  all_label = labels.float()
                  start_test = False
              else:
                  all_output = torch.cat((all_output, outputs.float().cpu()), 0)
                  all_label = torch.cat((all_label, labels.float()), 0)
      _, predict = torch.max(all_output, 1)
      all_output = nn.Softmax(dim=1)(all_output)
      ent = torch.sum(-all_output * torch.log(all_output + args.epsilon), dim=1) / np.log(all_output.size(1))
      accuracy = torch.sum(torch.squeeze(predict).float() == all_label).item() / float(all_label.size()[0])
      mean_ent = torch.mean(loss.Entropy(nn.Softmax(dim=1)(all_output))).cpu().data.item() 
      
      return accuracy, mean_ent
  
  def cal_acc_rot(loader, netF, netR):
      start_test = True
      with torch.no_grad():
          iter_test = iter(loader)
          for i in range(len(loader)):
              data = iter_test.next()
              inputs = data[0].cuda()
              r_labels = np.random.randint(0, 4, len(inputs))
              r_inputs = rotation.rotate_batch_with_labels(inputs, r_labels)
              r_labels = torch.from_numpy(r_labels)
              r_inputs = r_inputs.cuda()
  
              f_outputs, _ = netF(inputs)
              f_r_outputs, _ = netF(r_inputs)
              r_outputs = netR(torch.cat((f_outputs, f_r_outputs), 1))
  
              if start_test:
                  all_output = r_outputs.float().cpu()
                  all_label = r_labels.float()
                  start_test = False
              else:
                  all_output = torch.cat((all_output, r_outputs.float().cpu()), 0)
                  all_label = torch.cat((all_label, r_labels.float()), 0)
      _, predict = torch.max(all_output, 1)
      accuracy = torch.sum(torch.squeeze(predict).float() == all_label).item() / float(all_label.size()[0])
      
      return accuracy*100
  
  def train_target_rot(args):
      dset_loaders = data_load(args)
      netF = network.Res50().cuda()
      netR = network.feat_classifier(type='linear', class_num=4, bottleneck_dim=2*2048).cuda()
  
      netF.eval()
      for k, v in netF.named_parameters():
          v.requires_grad = False
  
      param_group = []
      for k, v in netR.named_parameters():
          param_group += [{'params': v, 'lr': args.lr*1}]
      netR.train()
      optimizer = optim.SGD(param_group)
      optimizer = op_copy(optimizer)
  
      max_iter = args.max_epoch * len(dset_loaders["target"])
      interval_iter = max_iter // 10
      iter_num = 0
  
      rot_acc = 0
      while iter_num < max_iter:
          optimizer.zero_grad()
          try:
              inputs_test, _, tar_idx = iter_test.next()
          except:
              iter_test = iter(dset_loaders["target"])
              inputs_test, _, tar_idx = iter_test.next()
  
          if inputs_test.size(0) == 1:
              continue
  
          inputs_test = inputs_test.cuda()
  
          iter_num += 1
          lr_scheduler(optimizer, iter_num=iter_num, max_iter=max_iter)
  
          r_labels_target = np.random.randint(0, 4, len(inputs_test))
          r_inputs_target = rotation.rotate_batch_with_labels(inputs_test, r_labels_target)
          r_labels_target = torch.from_numpy(r_labels_target).cuda()
          r_inputs_target = r_inputs_target.cuda()
          f_outputs, _ = netF(inputs_test)
          f_r_outputs, _ = netF(r_inputs_target)
          r_outputs_target = netR(torch.cat((f_outputs, f_r_outputs), 1))
          rotation_loss = nn.CrossEntropyLoss()(r_outputs_target, r_labels_target)
          rotation_loss.backward() 
          optimizer.step()
  
          if iter_num % interval_iter == 0 or iter_num == max_iter:
              netR.eval()
              acc_rot = cal_acc_rot(dset_loaders['test'], netF, netR)
              log_str = 'Task: {}, Iter:{}/{}; Accuracy = {:.2f}%'.format(args.name, iter_num, max_iter, acc_rot)
              args.out_file.write(log_str + '\n')
              args.out_file.flush()
              print(log_str+'\n')
              netR.train()
  
              if rot_acc < acc_rot:
                  rot_acc = acc_rot
                  best_netR = netR.state_dict()
  
      log_str = 'Best Accuracy = {:.2f}%'.format(rot_acc)
      args.out_file.write(log_str + '\n')
      args.out_file.flush()
      print(log_str+'\n')
  
      return best_netR, rot_acc
  
  def train_target(args):
      dset_loaders = data_load(args)  
      netF = network.Res50().cuda()
      if args.ssl > 0:
          netR = network.feat_classifier(type='linear', class_num=4, bottleneck_dim=2*2048).cuda()
          netR_dict, acc_rot = train_target_rot(args)
          netR.load_state_dict(netR_dict)
  
      param_group = []  
      for k, v in netF.named_parameters():
          if k.__contains__("fc"):
              v.requires_grad = False
          else:
              param_group += [{'params': v, 'lr': args.lr*args.lr_decay1}]
  
      if args.ssl > 0:
          for k, v in netR.named_parameters():
              param_group += [{'params': v, 'lr': args.lr*args.lr_decay2}]
          netR.train()
  
      optimizer = optim.SGD(param_group)
      optimizer = op_copy(optimizer)
  
      max_iter = args.max_epoch * len(dset_loaders["target"])
      interval_iter = max_iter // args.interval
      iter_num = 0
  
      netF.train()
      while iter_num < max_iter:
          try:
              inputs_test, _, tar_idx = iter_test.next()
          except:
              iter_test = iter(dset_loaders["target"])
              inputs_test, _, tar_idx = iter_test.next()
  
          if inputs_test.size(0) == 1:
              continue
  
          if iter_num % interval_iter == 0 and args.cls_par > 0:
              netF.eval()
              mem_label = obtain_label(dset_loaders['test'], netF, args)
              mem_label = torch.from_numpy(mem_label).cuda()
              netF.train()
  
          inputs_test = inputs_test.cuda()
          iter_num += 1
          lr_scheduler(optimizer, iter_num=iter_num, max_iter=max_iter)
  
          features_test, outputs_test = netF(inputs_test)
  
          if args.cls_par > 0:
              pred = mem_label[tar_idx]
              classifier_loss = nn.CrossEntropyLoss()(outputs_test, pred)
              classifier_loss *= args.cls_par
          else:
              classifier_loss = torch.tensor(0.0).cuda()
  
          if args.ent:
              softmax_out = nn.Softmax(dim=1)(outputs_test)
              entropy_loss = torch.mean(loss.Entropy(softmax_out))
              if args.gent:
                  msoftmax = softmax_out.mean(dim=0)
                  gentropy_loss = torch.sum(-msoftmax * torch.log(msoftmax + args.epsilon))
                  entropy_loss -= gentropy_loss
              classifier_loss += entropy_loss * args.ent_par
  
          optimizer.zero_grad()
          classifier_loss.backward()
  
          del features_test, outputs_test
          
          if args.ssl > 0:
              r_labels_target = np.random.randint(0, 4, len(inputs_test))
              r_inputs_target = rotation.rotate_batch_with_labels(inputs_test, r_labels_target)
              r_labels_target = torch.from_numpy(r_labels_target).cuda()
              r_inputs_target = r_inputs_target.cuda()
  
              f_outputs, _ = netF(inputs_test)
              f_outputs = f_outputs.detach()
  
              f_r_outputs, _ = netF(r_inputs_target)
              r_outputs_target = netR(torch.cat((f_outputs, f_r_outputs), 1))
  
              rotation_loss = args.ssl * nn.CrossEntropyLoss()(r_outputs_target, r_labels_target)   
              rotation_loss.backward() 
  
          optimizer.step()
  
          if iter_num % interval_iter == 0 or iter_num == max_iter:
              netF.eval()
              acc, ment = cal_acc(dset_loaders['test'], netF)
              log_str = 'Task: {}, Iter:{}/{}; Accuracy = {:.2f}%'.format(args.dset, iter_num, max_iter, acc*100)
              args.out_file.write(log_str + '\n')
              args.out_file.flush()
              print(log_str+'\n')
              netF.train()
      
      if args.issave:
          torch.save(netF.state_dict(), osp.join(args.output_dir, "target_" + args.savename + ".pt"))
  
      return netF
  
  def print_args(args):
      s = "==========================================\n"
      for arg, content in args.__dict__.items():
          s += "{}:{}\n".format(arg, content)
      return s
  
  def obtain_label(loader, net, args):
      start_test = True
      with torch.no_grad():
          iter_test = iter(loader)
          for _ in range(len(loader)):
              data = iter_test.next()
              inputs = data[0]
              labels = data[1]
              inputs = inputs.cuda()
              feas, outputs = net(inputs)
              if start_test:
                  all_fea = feas.float().cpu()
                  all_output = outputs.float().cpu()
                  all_label = labels.float()
                  start_test = False
              else:
                  all_fea = torch.cat((all_fea, feas.float().cpu()), 0)
                  all_output = torch.cat((all_output, outputs.float().cpu()), 0)
                  all_label = torch.cat((all_label, labels.float()), 0)
  
      all_output = nn.Softmax(dim=1)(all_output)
      ent = torch.sum(-all_output * torch.log(all_output + args.epsilon), dim=1)
      unknown_weight = 1 - ent / np.log(args.class_num)
      _, predict = torch.max(all_output, 1)
  
      accuracy = torch.sum(torch.squeeze(predict).float() == all_label).item() / float(all_label.size()[0])
      if args.distance == 'cosine':
          all_fea = torch.cat((all_fea, torch.ones(all_fea.size(0), 1)), 1)
          all_fea = (all_fea.t() / torch.norm(all_fea, p=2, dim=1)).t()
  
      all_fea = all_fea.float().cpu().numpy()
      K = all_output.size(1)
      aff = all_output.float().cpu().numpy()
      initc = aff.transpose().dot(all_fea)
      initc = initc / (1e-8 + aff.sum(axis=0)[:,None])
      cls_count = np.eye(K)[predict].sum(axis=0)
      labelset = np.where(cls_count>args.threshold)
      labelset = labelset[0]
  
      dd = cdist(all_fea, initc[labelset], args.distance)
      pred_label = dd.argmin(axis=1)
      pred_label = labelset[pred_label]
  
      for round in range(1):
          aff = np.eye(K)[pred_label]
          initc = aff.transpose().dot(all_fea)
          initc = initc / (1e-8 + aff.sum(axis=0)[:,None])
          dd = cdist(all_fea, initc[labelset], args.distance)
          pred_label = dd.argmin(axis=1)
          pred_label = labelset[pred_label]
  
      acc = np.sum(pred_label == all_label.float().numpy()) / len(all_fea)
      log_str = 'Accuracy = {:.2f}% -> {:.2f}%'.format(accuracy*100, acc*100)
  
      args.out_file.write(log_str + '\n')
      args.out_file.flush()
      print(log_str+'\n')
  
      return pred_label.astype('int') 
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser(description='SHOT++')
      parser.add_argument('--gpu_id', type=str, nargs='?', default='0', help="device id to run")
      parser.add_argument('--max_epoch', type=int, default=15, help="max iterations")
      parser.add_argument('--interval', type=int, default=15, help="max iterations")
      parser.add_argument('--batch_size', type=int, default=64, help="batch_size")
      parser.add_argument('--worker', type=int, default=4, help="number of workers")
      parser.add_argument('--distance', type=str, default='cosine', choices=["euclidean", "cosine"])
      parser.add_argument('--dset', type=str, default='imagenet_caltech', choices=['imagenet_caltech'])
      parser.add_argument('--lr', type=float, default=1e-2, help="learning rate")
      parser.add_argument('--net', type=str, default='resnet50', help="vgg16, resnet50")
      parser.add_argument('--seed', type=int, default=2019, help="random seed")
      parser.add_argument('--epsilon', type=float, default=1e-5)  
      parser.add_argument('--gent', type=bool, default=False)
      parser.add_argument('--ent', type=bool, default=True)
      parser.add_argument('--threshold', type=int, default=30)
  
      parser.add_argument('--cls_par', type=float, default=0.3)
      parser.add_argument('--ent_par', type=float, default=1.0)
      parser.add_argument('--output', type=str, default='seed')
      parser.add_argument('--da', type=str, default='pda', choices=['pda'])
      parser.add_argument('--ssl', type=float, default=0.0)
      parser.add_argument('--issave', type=bool, default=True)
      parser.add_argument('--lr_decay1', type=float, default=0.1)
      parser.add_argument('--lr_decay2', type=float, default=1.0)
  
      args = parser.parse_args()        
  
      os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
      SEED = args.seed
      torch.manual_seed(SEED)
      torch.cuda.manual_seed(SEED)
      np.random.seed(SEED)
      random.seed(SEED)
      # torch.backends.cudnn.deterministic = True
  
      args.class_num = 1000
      folder = '../data/'
      
      args.t_dset_path = folder + args.dset + '/' + 'caltech_84' + '_list.txt'
      args.test_dset_path = args.t_dset_path
  
      args.output_dir = osp.join(args.output, args.da, args.dset)
      args.name = args.dset
  
      if not osp.exists(args.output_dir):
          os.system('mkdir -p ' + args.output_dir)
      if not osp.exists(args.output_dir):
          os.mkdir(args.output_dir)
  
      args.savename = 'par_' + str(args.cls_par)
      if args.da == 'pda' and args.cls_par > 0:
          args.savename += ('_thr' + str(args.threshold))
      if args.ssl > 0:
          args.savename += ('_ssl_' + str(args.ssl))
      args.out_file = open(osp.join(args.output_dir, 'log_' + args.savename + '.txt'), 'w')
      args.out_file.write(print_args(args)+'\n')
      args.out_file.flush()
  
      train_target(args)
      
      args.out_file.close()