text: 
  # coding: utf-8
  """
  Pengyi Zhang
  201906
  """
  import cv2
  
  import argparse
  import json
  import os
  
  import numpy
  
  import torch
  import torch.nn as nn
  
  from torch.utils.data import DataLoader
  
  from models import *
  from utils.datasets import *
  from utils.utils import *
  from utils.parse_config import *
  
  """ Slim Principle
  (1) Use global threshold to control pruning ratio
  (2) Use local threshold to keep at least 10% unpruned 
  """
  
  def route_conv(layer_index, module_defs):
      """ find the convolutional layers connected by route layer
      """
      module_def = module_defs[layer_index]
      mtype = module_def['type']
      
      before_conv_id = []
      if mtype in ['convolutional', 'shortcut', 'upsample', 'maxpool']:
          if module_defs[layer_index-1]['type'] == 'convolutional':
              return [layer_index-1]
          before_conv_id += route_conv(layer_index-1, module_defs)
  
      elif mtype == "route":
          layer_is = [int(x)+layer_index if int(x) < 0 else int(x) for x in module_defs[layer_index]['layers'].split(',')]
          for layer_i in layer_is: 
              if module_defs[layer_i]['type'] == 'convolutional':
                  before_conv_id += [layer_i]
              else:
                  before_conv_id += route_conv(layer_i, module_defs)
          
      return before_conv_id
  
  
  def write_model_cfg(old_path, new_path, new_module_defs):
      """Parses the yolo-v3 layer configuration file and returns module definitions"""
      lines = []
      with open(old_path, 'r') as fp:
          old_lines = fp.readlines()
      for _line in old_lines:
          if "[convolutional]" in _line:
              break
          lines.append(_line)
  
      for i, module_def in enumerate(new_module_defs):
          
          mtype = module_def['type']
          lines.append("[{}]\n".format(mtype))
          print("layer:", i, mtype)
          if mtype == "convolutional":
              bn = 0
              filters = module_def['filters']
              bn = int(module_def['batch_normalize'])
              if bn:
                  lines.append("batch_normalize={}\n".format(bn))
                  filters = torch.sum(module_def['mask']).cpu().numpy().astype('int')        
              lines.append("filters={}\n".format(filters))
              lines.append("size={}\n".format(module_def['size']))
              lines.append("stride={}\n".format(module_def['stride']))
              lines.append("pad={}\n".format(module_def['pad']))
              lines.append("activation={}\n\n".format(module_def['activation']))
          elif mtype == "shortcut":
              lines.append("from={}\n".format(module_def['from']))
              lines.append("activation={}\n\n".format(module_def['activation']))   
          elif mtype == 'route':
              lines.append("layers={}\n\n".format(module_def['layers']))               
              
          elif mtype == 'upsample':
              lines.append("stride={}\n\n".format(module_def['stride']))
          elif mtype == 'maxpool':
              lines.append("stride={}\n".format(module_def['stride']))
              lines.append("size={}\n\n".format(module_def['size']))
          elif mtype == 'yolo':
              lines.append("mask = {}\n".format(module_def['mask']))
              lines.append("anchors = {}\n".format(module_def['anchors']))
              lines.append("classes = {}\n".format(module_def['classes']))
              lines.append("num = {}\n".format(module_def['num']))
              lines.append("jitter = {}\n".format(module_def['jitter']))
              lines.append("ignore_thresh = {}\n".format(module_def['ignore_thresh']))
              lines.append("truth_thresh = {}\n".format(module_def['truth_thresh']))
              lines.append("random = {}\n\n".format(module_def['random']))
      
      with open(new_path, "w") as f:
          f.writelines(lines)
  
  
          
  def test(
          cfg,
          weights=None,
          img_size=406,
          save=None,
          overall_ratio=0.5,
          perlayer_ratio=0.1
  ):
  
      """prune yolov3 and generate cfg, weights
      """
      if save != None:
          if not os.path.exists(save):
              os.makedirs(save)
      device = torch_utils.select_device()
      # Initialize model
      model = Darknet(cfg, img_size).to(device)
  
      # Load weights
      if weights.endswith('.pt'):  # pytorch format
          _state_dict = torch.load(weights, map_location=device)['model']
          model.load_state_dict(_state_dict)
      else:  # darknet format
          _ = load_darknet_weights(model, weights)
  
  ##  output a new cfg file
      total = 0
      for m in model.modules():
          if isinstance(m, nn.BatchNorm2d):
              total += m.weight.data.shape[0] # channels numbers
      
      bn = torch.zeros(total)
      index = 0
  
      for m in model.modules():
          if isinstance(m, nn.BatchNorm2d):
              size = m.weight.data.shape[0]
              bn[index:(index+size)] = m.weight.data.abs().clone()
              index += size
  
      sorted_bn, sorted_index = torch.sort(bn)
      thresh_index = int(total*overall_ratio)
      thresh = sorted_bn[thresh_index].cuda()
  
      print("--"*30)
      print()
      #print(list(model.modules()))
      # 
      proned_module_defs = model.module_defs
      for i, (module_def, module) in enumerate(zip(model.module_defs, model.module_list)):
          print("layer:", i)
          mtype = module_def['type']
          if mtype  == 'convolutional':
              bn = int(module_def['batch_normalize'])
              if bn:
                  m = getattr(module, 'batch_norm_%d' % i) # batch_norm layer
                  weight_copy = m.weight.data.abs().clone()
                  channels = weight_copy.shape[0] #
                  min_channel_num = int(channels * perlayer_ratio) if int(channels * perlayer_ratio) > 0 else 1
                  mask = weight_copy.gt(thresh).float().cuda()  
                  
                  if int(torch.sum(mask)) < min_channel_num: 
                      _, sorted_index_weights = torch.sort(weight_copy,descending=True)
                      mask[sorted_index_weights[:min_channel_num]]=1. 
  
                  proned_module_defs[i]['mask'] = mask.clone()
  
                  print('layer index: {:d} \t total channel: {:d} \t remaining channel: {:d}'.
                          format(i, mask.shape[0], int(torch.sum(mask)))) 
  
              print("layer:", mtype)
  
          elif mtype in ['upsample', 'maxpool']:
              print("layer:", mtype)
  
          elif mtype == 'route':
              print("layer:", mtype)
              # 
  
          elif mtype == 'shortcut':
              layer_i = int(module_def['from'])+i
              print("from layer ", layer_i)
              print("layer:", mtype)
              proned_module_defs[i]['is_access'] = False
              
  
          elif mtype == 'yolo':
              print("layer:", mtype)
              
  
      layer_number = len(proned_module_defs)
      for i in range(layer_number-1, -1, -1):
          mtype = proned_module_defs[i]['type']
          if mtype == 'shortcut': 
              if proned_module_defs[i]['is_access']: 
                  continue
  
              Merge_masks =  []
              layer_i = i
              while mtype == 'shortcut':
                  proned_module_defs[layer_i]['is_access'] = True
  
                  if proned_module_defs[layer_i-1]['type'] == 'convolutional': 
                      bn = int(proned_module_defs[layer_i-1]['batch_normalize'])
                      if bn: 
                          Merge_masks.append(proned_module_defs[layer_i-1]["mask"].unsqueeze(0))
  
                  layer_i = int(proned_module_defs[layer_i]['from'])+layer_i 
                  mtype = proned_module_defs[layer_i]['type']
  
                  if mtype == 'convolutional':              
                      bn = int(proned_module_defs[layer_i]['batch_normalize'])
                      if bn: 
                          Merge_masks.append(proned_module_defs[layer_i]["mask"].unsqueeze(0))
                  
  
              if len(Merge_masks) > 1:
                  Merge_masks = torch.cat(Merge_masks, 0)
                  merge_mask = (torch.sum(Merge_masks, dim=0) > 0).float().cuda()
              else:
                  merge_mask = Merge_masks[0].float().cuda()
  
              layer_i = i
              mtype = 'shortcut'
              while mtype == 'shortcut':
  
                  if proned_module_defs[layer_i-1]['type'] == 'convolutional': 
                      bn = int(proned_module_defs[layer_i-1]['batch_normalize'])
                      if bn:
                          proned_module_defs[layer_i-1]["mask"] = merge_mask
  
                  layer_i = int(proned_module_defs[layer_i]['from'])+layer_i 
                  mtype = proned_module_defs[layer_i]['type']
  
                  if mtype == 'convolutional': 
                      bn = int(proned_module_defs[layer_i]['batch_normalize'])
                      if bn:     
                          proned_module_defs[layer_i]["mask"] = merge_mask
  
  
  
      for i, (module_def, module) in enumerate(zip(model.module_defs, model.module_list)):
          print("layer:", i)
          mtype = module_def['type']
          if mtype  == 'convolutional':
              bn = int(module_def['batch_normalize'])
              if bn:
  
                  layer_i_1 = i - 1
                  proned_module_defs[i]['mask_before'] = None
  
                  mask_before = []
                  conv_indexs = []
                  if i > 0:
                      conv_indexs = route_conv(i, proned_module_defs)
                      for conv_index in conv_indexs:
                          mask_before += proned_module_defs[conv_index]["mask"].clone().cpu().numpy().tolist()
                      proned_module_defs[i]['mask_before'] = torch.tensor(mask_before).float().cuda()  
                     
                              
  
   
      output_cfg_path = os.path.join(save, "prune.cfg")
      write_model_cfg(cfg, output_cfg_path, proned_module_defs)
  
      pruned_model = Darknet(output_cfg_path, img_size).to(device)
      print(list(pruned_model.modules()))
      for i, (module_def, old_module, new_module) in enumerate(zip(proned_module_defs, model.module_list, pruned_model.module_list)):  
          mtype = module_def['type']
          print("layer: ",i, mtype)
          if mtype  == 'convolutional': # 
              bn = int(module_def['batch_normalize'])
              if bn:
                  new_norm = getattr(new_module, 'batch_norm_%d' % i) # batch_norm layer
                  old_norm = getattr(old_module, 'batch_norm_%d' % i) # batch_norm layer
  
                  new_conv = getattr(new_module, 'conv_%d' % i) # conv layer
                  old_conv = getattr(old_module, 'conv_%d' % i) # conv layer  
                  
  
                  idx1 = np.squeeze(np.argwhere(np.asarray(module_def['mask'].cpu().numpy())))
                  if i > 0:
                      idx2 = np.squeeze(np.argwhere(np.asarray(module_def['mask_before'].cpu().numpy())))
                      new_conv.weight.data = old_conv.weight.data[idx1.tolist()][:, idx2.tolist(), :, :].clone()
                      
                      print("idx1: ", len(idx1), ", idx2: ", len(idx2))
                  else:
                      new_conv.weight.data = old_conv.weight.data[idx1.tolist()].clone()
  
                  new_norm.weight.data = old_norm.weight.data[idx1.tolist()].clone()
                  new_norm.bias.data = old_norm.bias.data[idx1.tolist()].clone()
                  new_norm.running_mean = old_norm.running_mean[idx1.tolist()].clone()
                  new_norm.running_var = old_norm.running_var[idx1.tolist()].clone()
                  
  
                  print('layer index: ', i, 'idx1: ', idx1)     
              else: 
  
                  new_conv = getattr(new_module, 'conv_%d' % i) # batch_norm layer
                  old_conv = getattr(old_module, 'conv_%d' % i) # batch_norm layer
                  idx2 = np.squeeze(np.argwhere(np.asarray(proned_module_defs[i-1]['mask'].cpu().numpy())))
                  new_conv.weight.data = old_conv.weight.data[:,idx2.tolist(),:,:].clone()
                  new_conv.bias.data = old_conv.bias.data.clone()
                  print('layer index: ', i, "entire copy") 
  
      print('--'*30)
      print('prune done!')    
      print('pruned ratio %.3f'%overall_ratio)
      prune_weights_path = os.path.join(save, "prune.pt")    
      _pruned_state_dict = pruned_model.state_dict()
      torch.save(_pruned_state_dict, prune_weights_path)
  
      print("Done!") 
  
  
  
      # test
      pruned_model.eval()
      img_path = "test.jpg"
      
      org_img = cv2.imread(img_path)  # BGR
      img, ratiow, ratioh, padw, padh = letterbox(org_img, new_shape=[img_size,img_size], mode='rect')
  
      # Normalize
      img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
      img = np.ascontiguousarray(img, dtype=np.float32)  # uint8 to float32
      img /= 255.0  # 0 - 255 to 0.0 - 1.0
  
      imgs = torch.from_numpy(img).unsqueeze(0).to(device)
      _, _, height, width = imgs.shape  # batch size, channels, height, width
  
      # Run model
      inf_out, train_out = pruned_model(imgs)  # inference and training outputs
      # Run NMS
      output = non_max_suppression(inf_out, conf_thres=0.005, nms_thres=0.5)
      # Statistics per image
      for si, pred in enumerate(output):
          if pred is None:
              continue
          if True:
              box = pred[:, :4].clone()  # xyxy
              scale_coords(imgs[si].shape[1:], box, org_img.shape[:2])  # to original shape
              for di, d in enumerate(pred):
                  category_id = int(d[6])
                  left, top, right, bot = [float(x) for x in box[di]]
                  confidence = float(d[4])
  
                  cv2.rectangle(org_img, (int(left), int(top)), (int(right), int(bot)),
                                  (255, 0, 0), 2)
                  cv2.putText(org_img, str(category_id) + ":" + str('%.1f' % (float(confidence) * 100)) + "%", (int(left), int(top) - 8),
                              cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1)  
          cv2.imshow("result", org_img)
          cv2.waitKey(-1)            
          cv2.imwrite('result_{}'.format(img_path), org_img)
  
  
      # convert pt to weights:
      prune_c_weights_path = os.path.join(save, "prune.weights")
      save_weights(pruned_model, prune_c_weights_path)
      
  
  if __name__ == '__main__':
      os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
      os.environ["CUDA_VISIBLE_DEVICES"] = "0"
      parser = argparse.ArgumentParser(description='PyTorch Slimming Yolov3 prune')
      parser.add_argument('--cfg', type=str, default='VisDrone2019/yolov3-spp3.cfg', help='cfg file path')
      parser.add_argument('--weights', type=str, default='yolov3-spp3_final.weights', help='path to weights file')
      parser.add_argument('--img_size', type=int, default=608, help='inference size (pixels)')
      parser.add_argument('--save', default='prune', type=str, metavar='PATH', help='path to save pruned model (default: none)')
      parser.add_argument('--overall_ratio', type=float, default=0.5, help='scale sparse rate (default: 0.5)')    
      parser.add_argument('--perlayer_ratio', type=float, default=0.1, help='minimal scale sparse rate (default: 0.1) to prevent disconnect')    
      
      opt = parser.parse_args()
      opt.save += "_{}_{}".format(opt.overall_ratio, opt.perlayer_ratio)
  
      print(opt)
  
      with torch.no_grad():
          test(
              opt.cfg,
              opt.weights,
              opt.img_size,
              opt.save,
              opt.overall_ratio,
              opt.perlayer_ratio,
          )
  