text: # -*- coding: utf-8 -*-
  
  from __future__ import print_function, division
  
  import argparse
  import torch
  import torch.nn as nn
  import torch.optim as optim
  from torch.optim import lr_scheduler
  from torch.autograd import Variable
  import numpy as np
  import torchvision
  from torchvision import datasets, models, transforms
  import matplotlib
  matplotlib.use('agg')
  import matplotlib.pyplot as plt
  from PIL import Image
  import time
  import os
  from model import ft_net, ft_net_dense, PCB
  from random_erasing import RandomErasing
  import json
  from shutil import copyfile
  
  version =  torch.__version__
  
  #PN GAN code integration
  from torch.autograd import Variable
  import network
  import dataset
  from config import cfg
  
  # Load Network
  def load_network(model_path):
      os.environ["CUDA_VISIBLE_DEVICES"] = cfg.TEST.GPU_ID
      print ('###################################')
      print ("#####      Load Network      #####")
      print ('###################################')
  
      nets = []
      netG = network.Res_Generator(cfg.TRAIN.ngf, cfg.TRAIN.num_resblock)
      netG.load_state_dict(torch.load(model_path)['state_dict'])
  
      nets.append(netG)
      for net in nets:
          net.cuda()
  
      print ('Finished !')
      return nets
  
  ######################################################################
  # Options
  # --------
  parser = argparse.ArgumentParser(description='Training')
  parser.add_argument('--gpu_ids',default='0', type=str,help='gpu_ids: e.g. 0  0,1,2  0,2')
  parser.add_argument('--name',default='ft_ResNet50', type=str, help='output model name')
  parser.add_argument('--data_dir',default='/home/zzd/Market/pytorch',type=str, help='training dir path')
  parser.add_argument('--train_all', action='store_true', help='use all training data' )
  parser.add_argument('--color_jitter', action='store_true', help='use color jitter in training' )
  parser.add_argument('--batchsize', default=32, type=int, help='batchsize')
  parser.add_argument('--erasing_p', default=0, type=float, help='Random Erasing probability, in [0,1]')
  parser.add_argument('--use_dense', action='store_true', help='use densenet121' )
  parser.add_argument('--PCB', action='store_true', help='use PCB+ResNet50' )
  parser.add_argument('--use_gan', action='store_true', help="Use pose normalization or not")
  opt = parser.parse_args()
  
  global pn_gan
  global poses
      
  if args.use_gan:
          pn_gan,poses = initialize_PN_GAN();
          
  pn_gan=''
  poses =''
  
  def initialize_PN_GAN():
      model_path = './G_2.pkl'
      pn_gan = load_network(model_path)
      pose_data = dataset.Pose_Loader(pose_path='./poses/',transform=dataset.val_transform(), loader=dataset.val_loader)
      use_gpu = torch.cuda.is_available()
      pin_memory = True if use_gpu else False    
      return pn_gan[0], DataLoader(pose_data,batch_size=1,shuffle=False,num_workers=1,pin_memory=pin_memory)    
  
  
  data_dir = opt.data_dir
  name = opt.name
  str_ids = opt.gpu_ids.split(',')
  gpu_ids = []
  for str_id in str_ids:
      gid = int(str_id)
      if gid >=0:
          gpu_ids.append(gid)
  
  # set gpu ids
  if len(gpu_ids)>0:
      torch.cuda.set_device(gpu_ids[0])
  #print(gpu_ids[0])
  
  
  ######################################################################
  # Load Data
  # ---------
  #
  
  transform_train_list = [
          #transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC)
          transforms.Resize((288,144), interpolation=3),
          transforms.RandomCrop((256,128)),
          transforms.RandomHorizontalFlip(),
          transforms.ToTensor(),
          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
          ]
  
  transform_val_list = [
          transforms.Resize(size=(256,128),interpolation=3), #Image.BICUBIC
          transforms.ToTensor(),
          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
          ]
  
  if opt.PCB:
      transform_train_list = [
          transforms.Resize((384,192), interpolation=3),
          transforms.RandomHorizontalFlip(),
          transforms.ToTensor(),
          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
          ]
      transform_val_list = [
          transforms.Resize(size=(384,192),interpolation=3), #Image.BICUBIC
          transforms.ToTensor(),
          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
          ]
  
  if opt.erasing_p>0:
      transform_train_list = transform_train_list +  [RandomErasing(probability = opt.erasing_p, mean=[0.0, 0.0, 0.0])]
  
  if opt.color_jitter:
      transform_train_list = [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0)] + transform_train_list
  
  print(transform_train_list)
  data_transforms = {
      'train': transforms.Compose( transform_train_list ),
      'val': transforms.Compose(transform_val_list),
  }
  
  
  train_all = ''
  if opt.train_all:
       train_all = '_all'
  
  image_datasets = {}
  image_datasets['train'] = datasets.ImageFolder(os.path.join(data_dir, 'train' + train_all),
                                            data_transforms['train'])
  image_datasets['val'] = datasets.ImageFolder(os.path.join(data_dir, 'val'),
                                            data_transforms['val'])
  
  dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=opt.batchsize,
                                               shuffle=True, num_workers=8) # 8 workers may work faster
                for x in ['train', 'val']}
  dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
  class_names = image_datasets['train'].classes
  
  use_gpu = torch.cuda.is_available()
  
  since = time.time()
  inputs, classes = next(iter(dataloaders['train']))
  print(time.time()-since)
  ######################################################################
  # Training the model
  # ------------------
  #
  # Now, let's write a general function to train a model. Here, we will
  # illustrate:
  #
  # -  Scheduling the learning rate
  # -  Saving the best model
  #
  # In the following, parameter ``scheduler`` is an LR scheduler object from
  # ``torch.optim.lr_scheduler``.
  
  y_loss = {} # loss history
  y_loss['train'] = []
  y_loss['val'] = []
  y_err = {}
  y_err['train'] = []
  y_err['val'] = []
  
  def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
      since = time.time()
  
      best_model_wts = model.state_dict()
      best_acc = 0.0
  
      for epoch in range(num_epochs):
          print('Epoch {}/{}'.format(epoch, num_epochs - 1))
          print('-' * 10)
  
          # Each epoch has a training and validation phase
          for phase in ['train', 'val']:
              if phase == 'train':
                  scheduler.step()
                  model.train(True)  # Set model to training mode
              else:
                  model.train(False)  # Set model to evaluate mode
  
              running_loss = 0.0
              running_corrects = 0.0
              # Iterate over data.
              for data in dataloaders[phase]:
                  # get the inputs
                  inputs, labels = data
                  now_batch_size,c,h,w = inputs.shape
                  if now_batch_size<opt.batchsize: # skip the last batch
                      continue
                  #print(inputs.shape)
                  # wrap them in Variable
                  if use_gpu:
                      inputs = Variable(inputs.cuda())
                      labels = Variable(labels.cuda())
                  else:
                      inputs, labels = Variable(inputs), Variable(labels)
  
                  # zero the parameter gradients
                  optimizer.zero_grad()
                  
                  if args.use_gan:
                      for pose_idx, pose_img in enumerate(poses):
                          if use_gpu:
                              pose_img = Variable(pose_img).cuda()
                          tgt_img = pn_gan(Variable(inputs).cuda(),pose_img)
                          
                          outputs = model(tgt_img)
                          if not opt.PCB:
                              _, preds = torch.max(outputs.data, 1)
                              loss = criterion(outputs, labels)
                          else:
                              part = {}
                              sm = nn.Softmax(dim=1)
                              num_part = 6
                              for i in range(num_part):
                                  part[i] = outputs[i]
  
                              score = sm(part[0]) + sm(part[1]) +sm(part[2]) + sm(part[3]) +sm(part[4]) +sm(part[5])
                              _, preds = torch.max(score.data, 1)
  
                              loss = criterion(part[0], labels)
                              for i in range(num_part-1):
                                  loss += criterion(part[i+1], labels)
  
                          # backward + optimize only if in training phase
                          if phase == 'train':
                              loss.backward()
                              optimizer.step()
                          
                  
                  # forward
                  outputs = model(inputs)
                  if not opt.PCB:
                      _, preds = torch.max(outputs.data, 1)
                      loss = criterion(outputs, labels)
                  else:
                      part = {}
                      sm = nn.Softmax(dim=1)
                      num_part = 6
                      for i in range(num_part):
                          part[i] = outputs[i]
  
                      score = sm(part[0]) + sm(part[1]) +sm(part[2]) + sm(part[3]) +sm(part[4]) +sm(part[5])
                      _, preds = torch.max(score.data, 1)
  
                      loss = criterion(part[0], labels)
                      for i in range(num_part-1):
                          loss += criterion(part[i+1], labels)
  
                  # backward + optimize only if in training phase
                  if phase == 'train':
                      loss.backward()
                      optimizer.step()
  
                  # statistics
                  if int(version[2]) > 3: # for the new version like 0.4.0 and 0.5.0
                      running_loss += loss.item() * now_batch_size
                  else :  # for the old version like 0.3.0 and 0.3.1
                      running_loss += loss.data[0] * now_batch_size
                  running_corrects += float(torch.sum(preds == labels.data))
  
              epoch_loss = running_loss / dataset_sizes[phase]
              epoch_acc = running_corrects / dataset_sizes[phase]
              
              print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                  phase, epoch_loss, epoch_acc))
              
              y_loss[phase].append(epoch_loss)
              y_err[phase].append(1.0-epoch_acc)            
              # deep copy the model
              if phase == 'val':
                  last_model_wts = model.state_dict()
                  if epoch%10 == 9:
                      save_network(model, epoch)
                  draw_curve(epoch)
  
          print()
  
      time_elapsed = time.time() - since
      print('Training complete in {:.0f}m {:.0f}s'.format(
          time_elapsed // 60, time_elapsed % 60))
      #print('Best val Acc: {:4f}'.format(best_acc))
  
      # load best model weights
      model.load_state_dict(last_model_wts)
      save_network(model, 'last')
      return model
  
  
  ######################################################################
  # Draw Curve
  #---------------------------
  x_epoch = []
  fig = plt.figure()
  ax0 = fig.add_subplot(121, title="loss")
  ax1 = fig.add_subplot(122, title="top1err")
  def draw_curve(current_epoch):
      x_epoch.append(current_epoch)
      ax0.plot(x_epoch, y_loss['train'], 'bo-', label='train')
      ax0.plot(x_epoch, y_loss['val'], 'ro-', label='val')
      ax1.plot(x_epoch, y_err['train'], 'bo-', label='train')
      ax1.plot(x_epoch, y_err['val'], 'ro-', label='val')
      if current_epoch == 0:
          ax0.legend()
          ax1.legend()
      fig.savefig( os.path.join('./model',name,'train.jpg'))
  
  ######################################################################
  # Save model
  #---------------------------
  def save_network(network, epoch_label):
      save_filename = 'net_%s.pth'% epoch_label
      save_path = os.path.join('./model',name,save_filename)
      torch.save(network.cpu().state_dict(), save_path)
      if torch.cuda.is_available():
          network.cuda(gpu_ids[0])
  
  
  ######################################################################
  # Finetuning the convnet
  # ----------------------
  #
  # Load a pretrainied model and reset final fully connected layer.
  #
  
  if opt.use_dense:
      model = ft_net_dense(len(class_names))
  else:
      model = ft_net(len(class_names))
  
  if opt.PCB:
      model = PCB(len(class_names))
  
  print(model)
  
  if use_gpu:
      model = model.cuda()
  
  criterion = nn.CrossEntropyLoss()
  
  if not opt.PCB:
      ignored_params = list(map(id, model.model.fc.parameters() )) + list(map(id, model.classifier.parameters() ))
      base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
      optimizer_ft = optim.SGD([
               {'params': base_params, 'lr': 0.01},
               {'params': model.model.fc.parameters(), 'lr': 0.1},
               {'params': model.classifier.parameters(), 'lr': 0.1}
           ], weight_decay=5e-4, momentum=0.9, nesterov=True)
  else:
      ignored_params = list(map(id, model.model.fc.parameters() ))
      ignored_params += (list(map(id, model.classifier0.parameters() )) 
                       +list(map(id, model.classifier1.parameters() ))
                       +list(map(id, model.classifier2.parameters() ))
                       +list(map(id, model.classifier3.parameters() ))
                       +list(map(id, model.classifier4.parameters() ))
                       +list(map(id, model.classifier5.parameters() ))
                       #+list(map(id, model.classifier6.parameters() ))
                       #+list(map(id, model.classifier7.parameters() ))
                        )
      base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
      optimizer_ft = optim.SGD([
               {'params': base_params, 'lr': 0.01},
               {'params': model.model.fc.parameters(), 'lr': 0.1},
               {'params': model.classifier0.parameters(), 'lr': 0.1},
               {'params': model.classifier1.parameters(), 'lr': 0.1},
               {'params': model.classifier2.parameters(), 'lr': 0.1},
               {'params': model.classifier3.parameters(), 'lr': 0.1},
               {'params': model.classifier4.parameters(), 'lr': 0.1},
               {'params': model.classifier5.parameters(), 'lr': 0.1},
               #{'params': model.classifier6.parameters(), 'lr': 0.01},
               #{'params': model.classifier7.parameters(), 'lr': 0.01}
           ], weight_decay=5e-4, momentum=0.9, nesterov=True)
  
  # Decay LR by a factor of 0.1 every 40 epochs
  exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=40, gamma=0.1)
  
  ######################################################################
  # Train and evaluate
  # ^^^^^^^^^^^^^^^^^^
  #
  # It should take around 1-2 hours on GPU. 
  #
  dir_name = os.path.join('./model',name)
  if not os.path.isdir(dir_name):
      os.mkdir(dir_name)
  #record every run
  copyfile('./train.py', dir_name+'/train.py')
  copyfile('./model.py', dir_name+'/model.py')
  
  # save opts
  with open('%s/opts.json'%dir_name,'w') as fp:
      json.dump(vars(opt), fp, indent=1)
  
  model = train_model(model, criterion, optimizer_ft, exp_lr_scheduler,
                         num_epochs=60)
  
  