text: import numpy as np
  import torch
  from torch.backends import cudnn
  cudnn.enabled = True
  from torch.utils.data import DataLoader, Dataset
  from torchvision import transforms
  import torch.optim as optim
  import voc12.data
  from tool import pyutils, imutils, torchutils
  import argparse
  import importlib
  import torch.nn.functional as F
  
  import matplotlib.pyplot as plt
  import math
  from tqdm import tqdm
  import collections
  
  import random
  import scipy.misc
  import os
  from PIL import Image
  from tensorboardX import SummaryWriter
  
  
  
  def compute_acc(pred_labels, gt_labels):
      pred_correct_count = 0
      pred_correct_list = []
      for pred_label in pred_labels:
          if pred_label in gt_labels:
              pred_correct_count += 1
      union = len(gt_labels) + len(pred_labels) - pred_correct_count
      acc = round(pred_correct_count/union, 4)
      return acc
  
  
  def cls200_sum(y_200):
      cls20_prob_sum_list = []
      for rou in range(20):
          subclass_prob_sum = sum(y_200[rou*k_cluster:rou*k_cluster+k_cluster])
          cls20_prob_sum_list.append(subclass_prob_sum/10)
  
      cls200_pred_max = np.where(np.array(cls20_prob_sum_list)>0.05)[0]
      return cls200_pred_max
  
  
  def get_img_path(img_name, dataset_path):
      tmp = os.path.join(dataset_path, img_name + '.jpg')
      return tmp
  
  
  
  
  
  class Sub_Class_Dataset(Dataset):
      def __init__(self, voc12_root, crop_size, round_nb, k_cluster, save_folder, test=False):
          print('############################################## || k_{} / Round {} || ##############################################'.format(k_cluster, round_nb))
  
          self.voc12_root = voc12_root
          self.crop_size = crop_size
  
          with open('{}/label/R{}_train_filename_list.txt'.format(save_folder, round_nb), 'r') as f:
              self.filename = f.read().split('\n')
          f.close()
          self.filename = self.filename[:-1]  # 16458
  
          self.label = np.load('{}/label/R{}_train_label_200.npy'.format(save_folder, round_nb), allow_pickle=True)
          self.label = torch.from_numpy(self.label).float()
  
          self.label_20 = np.load('{}/label/R{}_train_label_20.npy'.format(save_folder, round_nb), allow_pickle=True) # 16458
  
          self.label_20 = torch.from_numpy(self.label_20).float()
  
          print('=='*60)
          print('Training Data: image: {} | 20 class label: {} | 200 class label: {}'.format(len(self.filename), self.label_20.shape, self.label.shape))
          print('=='*60)
  
  
      def __getitem__(self, index):
  
          label_200 = self.label[index]
          label_20 = self.label_20[index]
  
          self.dataset_path = os.path.join(self.voc12_root, 'JPEGImages')
  
          filename = self.filename[index]
  
          img = Image.open(get_img_path(filename, self.dataset_path)).convert("RGB")
          img = imutils.ResizeLong(img, 256, 512)
          img = imutils.Flip(img)
          img = imutils.ColorJitter(img)
          img = np.array(img)
          img = imutils.NNormalize(img)
          img = imutils.Crop(img, self.crop_size)
          img = img.transpose(2,0,1)
          img = torch.from_numpy(img)
  
          return img, label_20, label_200, filename
  
      def __len__(self):
          return len(self.filename)
  
  
  
  
  if __name__ == '__main__':
  
      parser = argparse.ArgumentParser()
      parser.add_argument("--batch_size", default=16, type=int)
      parser.add_argument("--max_epoches", default=61, type=int)
      parser.add_argument("--network", default="network.resnet38_cls", type=str)
      parser.add_argument("--lr", default=0.01, type=float)
      parser.add_argument("--num_workers", default=8, type=int)
      parser.add_argument("--wt_dec", default=5e-4, type=float)
      parser.add_argument("--weights", required=True, type=str, help="the path to the pretrained weight ")
      parser.add_argument("--train_list", default="voc12/train_aug.txt", type=str)
      parser.add_argument("--session_name", default="resnet_cls", type=str)
      parser.add_argument("--crop_size", default=224, type=int)
      parser.add_argument("--voc12_root", default="/home/julia/datasets/VOC2012", type=str, help="the path to the dataset folder")
      parser.add_argument("--subcls_loss_weight", default="5", type=float, help="the weight multiply to the sub-category loss")
      parser.add_argument("--round_nb", default="0", type=int, help="the round number of the training classifier, e.g., 1st round: round_nb=1, and so on")
      parser.add_argument("--k_cluster", default="10", type=int, help="the number of the sub-category")
      parser.add_argument("--save_folder", required=True, default="./save", type=str, help="the path to save the model")
  
      args = parser.parse_args()
  
      model = getattr(importlib.import_module(args.network), 'Net')(args.k_cluster, args.round_nb)
  
      pyutils.Logger(args.session_name + '.log')
  
      print(vars(args))
  
      log_path = os.path.join(args.save_folder, 'log', 'R{}'.format(args.round_nb))
      writer = SummaryWriter('{}'.format(log_path))
  
  
      train_dataset = Sub_Class_Dataset(voc12_root=args.voc12_root,
                                          crop_size = args.crop_size,
                                          round_nb=args.round_nb,
                                          k_cluster=args.k_cluster,
                                          save_folder=args.save_folder)
  
      train_data_loader = DataLoader(train_dataset,
                                      batch_size=args.batch_size,
                                      shuffle=True,
                                      num_workers=args.num_workers,
                                      pin_memory=True,
                                      drop_last=True)
  
  
  
      max_step = (len(train_dataset) // args.batch_size) * args.max_epoches
  
  
      param_groups = model.get_parameter_groups()
      optimizer = torchutils.PolyOptimizer([
          {'params': param_groups[0], 'lr': args.lr, 'weight_decay': args.wt_dec},
          {'params': param_groups[1], 'lr': 2*args.lr, 'weight_decay': 0},
          {'params': param_groups[2], 'lr': 10*args.lr, 'weight_decay': args.wt_dec},
          {'params': param_groups[3], 'lr': 20*args.lr, 'weight_decay': 0}
      ], lr=args.lr, weight_decay=args.wt_dec, max_step=max_step)
  
  
      if args.weights[-7:] == '.params':
          assert args.network == "network.resnet38_cls"
          import network.resnet38d
          weights_dict = network.resnet38d.convert_mxnet_to_torch(args.weights)
      elif args.weights[-11:] == '.caffemodel':
          assert args.network == "network.vgg16_cls"
          import network.vgg16d
          weights_dict = network.vgg16d.convert_caffe_to_torch(args.weights)
      else:
          weights_dict = torch.load(args.weights)
  
      model.load_state_dict(weights_dict, strict=False)
      model = torch.nn.DataParallel(model).cuda()
      model.train()
  
      avg_meter = pyutils.AverageMeter('loss')
  
      timer = pyutils.Timer("Session started: ")
  
  
      parent_labels = np.load('./voc12/cls_labels.npy', allow_pickle=True).tolist()
      k_cluster = 10
  
  
      step = 0
      for ep in range(args.max_epoches):
          ep_count = 0
          ep_EM = 0
          ep_acc = 0
  
          ep_p_EM = 0
          ep_p_acc = 0
          ep_acc_vote = 0
  
          cls20_ep_EM = 0
          cls20_ep_acc = 0
  
          for iter, (data, label_20, label_200, filename) in tqdm(enumerate(train_data_loader)):
  
              img = data
              label_20 = label_20.cuda(non_blocking=True)
              label_200 = label_200.cuda(non_blocking=True)
              img_name = filename
  
  
              x_20, _, y_20, x_200, y_200 = model(img, args.round_nb)
  
              # compute acc for 20 classes
              cls20_prob = y_20.cpu().data.numpy()
              cls20_gt = label_20.cpu().data.numpy()
              for num, one in enumerate(cls20_prob):
                  ep_count += 1
                  pass_cls = np.where(one > 0.5)[0]
                  true_cls_20 = np.where(cls20_gt[num] == 1)[0]
  
                  if np.array_equal(pass_cls, true_cls_20) == True:  # exact match
                      cls20_ep_EM += 1
  
                  acc = compute_acc(pass_cls, true_cls_20)
                  cls20_ep_acc += acc
  
  
              # compute acc for 200 classes
              tmp = y_200.cpu().data.numpy()
              tmp_label = label_200.cpu().data.numpy()
              for num, one in enumerate(tmp):
                  pass_cls = np.where(one>0.5)[0]
                  true_cls_200 = np.where(tmp_label[num] == 1)[0]
  
                  parent_cls = np.where(parent_labels[img_name[num]] == 1)[0]
  
                  cls200_pred_max = cls200_sum(one)
  
                  if np.array_equal(pass_cls, true_cls_200) == True:
                      ep_EM += 1
  
                  # cls200: acc for 200-->20 (sum)
                  acc = compute_acc(cls200_pred_max, parent_cls)
                  ep_acc += acc
  
                  # cls200: acc for 200-->20 (top 1)
                  pass_map_cls = np.unique([int(m/k_cluster) for m in np.where(one > 0.5)[0]])
                  if np.array_equal(pass_map_cls, parent_cls) == True:
                      ep_p_EM += 1
  
                  p_acc = compute_acc(pass_map_cls, parent_cls)
                  ep_p_acc += p_acc
  
  
  
              avg_ep_EM = round(ep_EM/ep_count, 4)
              avg_ep_acc = round(ep_acc/ep_count, 4)
  
              avg_ep_p_EM = round(ep_p_EM/ep_count, 4)
              avg_ep_p_acc = round(ep_p_acc/ep_count, 4)
  
              avg_cls20_ep_EM = round(cls20_ep_EM/ep_count, 4)
              avg_cls20_ep_acc = round(cls20_ep_acc/ep_count, 4)
  
              cls_20_loss = F.multilabel_soft_margin_loss(x_20, label_20)
              cls_200_loss = F.multilabel_soft_margin_loss(x_200, label_200)
  
              loss = cls_20_loss + (args.subcls_loss_weight*cls_200_loss)
  
              if iter%100 ==0:
                  print('k{} R{}| Ep:{} L:{} -20_LOSS:{} -200_LOSS:{} | -cls20:{} | -cls200_Top1:{} -Sum:{}'.format(args.k_cluster, args.round_nb, ep, round(loss.item(), 3), round(cls_20_loss.item(), 3), round(cls_200_loss.item(), 3), avg_cls20_ep_acc, avg_ep_p_acc, avg_ep_acc))
  
              avg_meter.add({'loss': cls_200_loss.item()})
  
  
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
              torch.cuda.empty_cache()
  
              writer.add_scalar('20 Classification Loss', cls_20_loss.item(), step)
              writer.add_scalar('200 Classification Loss', cls_200_loss.item(), step)
              writer.add_scalar('Total Loss', loss.item(), step)
              writer.add_scalar('Cls20 Accuracy', avg_cls20_ep_acc, step)
              writer.add_scalar('Cls200 Accuracy (Sum)', avg_ep_acc, step)
              writer.add_scalar('Cls200 Accuracy (Top1)', avg_ep_p_acc, step)
  
              step += 1
  
              if (optimizer.global_step-1)%100 == 0:
                  timer.update_progress(optimizer.global_step / max_step)
  
                  print('Iter:%5d/%5d' % (optimizer.global_step - 1, max_step),
                        'Loss:%.4f' % (avg_meter.pop('loss')),
                        'imps:%.1f' % ((iter+1) * args.batch_size / timer.get_stage_elapsed()),
                        'Fin:%s' % (timer.str_est_finish()),
                        'lr: %.4f' % (optimizer.param_groups[0]['lr']), flush=True)
  
  
          if ep % 10 == 0:
              torch.save(model.module.state_dict(), '{}/weight/k{}_R{}_'.format(args.save_folder, args.k_cluster, args.round_nb) + args.session_name + '_ep{}.pth'.format(ep))
              print('Loss: {} achieves the lowest one => Epoch {} weights are saved!'.format(loss, ep))
  