scripts/train_segmentation_main.py

   1 #!/usr/bin/env python
   2
   3 import sys
   4 import os
   5 import cv2
   6 import argparse
   7 import datetime
   8 import numpy as np
   9
  10 ################################
  11 from pytorch_jacinto_ai.xnn.utils import str2bool
  12 parser = argparse.ArgumentParser()
  13 parser.add_argument('--save_path', type=str, default=None, help='checkpoint save folder')
  14 parser.add_argument('--gpus', type=int, nargs='*', default=None, help='Base learning rate')
  15 parser.add_argument('--batch_size', type=int, default=None, help='Batch size')
  16 parser.add_argument('--lr', type=float, default=None, help='Base learning rate')
  17 parser.add_argument('--lr_clips', type=float, default=None, help='Learning rate for clips in PAct2')
  18 parser.add_argument('--lr_calib', type=float, default=None, help='Learning rate for calibration')
  19 parser.add_argument('--model_name', type=str, default=None, help='model name')
  20 parser.add_argument('--dataset_name', type=str, default=None, help='dataset name')
  21 parser.add_argument('--data_path', type=str, default=None, help='data path')
  22 parser.add_argument('--epoch_size', type=float, default=None, help='epoch size. using a fraction will reduce the data used for one epoch')
  23 parser.add_argument('--epochs', type=int, default=None, help='number of epochs')
  24 parser.add_argument('--warmup_epochs', type=int, default=None, help='number of epochs for the learning rate to increase and reach base value')
  25 parser.add_argument('--milestones', type=int, nargs='*', default=None, help='change lr at these milestones')
  26 parser.add_argument('--img_resize', type=int, nargs=2, default=None, help='img_resize size. for training this will be modified according to rand_scale')
  27 parser.add_argument('--rand_scale', type=float, nargs=2, default=None, help='random scale factors for training')
  28 parser.add_argument('--rand_crop', type=int, nargs=2, default=None, help='random crop for training')
  29 parser.add_argument('--output_size', type=int, nargs=2, default=None, help='output size of the evaluation - prediction/groundtruth. this is not used while training as it blows up memory requirement')
  30 parser.add_argument('--pretrained', type=str, default=None, help='pretrained model')
  31 parser.add_argument('--resume', type=str, default=None, help='resume an unfinished training from this model')
  32 parser.add_argument('--phase', type=str, default=None, help='training/calibration/validation')
  33 parser.add_argument('--evaluate_start', type=str2bool, default=None, help='Whether to run validation before the training')
  34 #
  35 parser.add_argument('--quantize', type=str2bool, default=None, help='Quantize the model')
  36 parser.add_argument('--histogram_range', type=str2bool, default=None, help='run only evaluation and no training')
  37 parser.add_argument('--per_channel_q', type=str2bool, default=None, help='run only evaluation and no training')
  38 parser.add_argument('--bias_calibration', type=str2bool, default=None, help='run only evaluation and no training')
  39 parser.add_argument('--bitwidth_weights', type=int, default=None, help='bitwidth for weight quantization')
  40 parser.add_argument('--bitwidth_activations', type=int, default=None, help='bitwidth for activation quantization')
  41 #
  42 parser.add_argument('--freeze_bn', type=str2bool, default=None, help='freeze the bn stats or not')
  43 cmds = parser.parse_args()
  44
  45 ################################
  46 # taken care first, since this has to be done before importing pytorch
  47 if 'gpus' in vars(cmds):
  48     value = getattr(cmds, 'gpus')
  49     if (value is not None) and ("CUDA_VISIBLE_DEVICES" not in os.environ):
  50         os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(v) for v in value])
  51     #
  52 #
  53
  54 ################################
  55 # to avoid hangs in data loader with multi threads
  56 # this was observed after using cv2 image processing functions
  57 # https://github.com/pytorch/pytorch/issues/1355
  58 cv2.setNumThreads(0)
  59
  60 ################################
  61 #import of torch should be after CUDA_VISIBLE_DEVICES for it to take effect
  62 import torch
  63 from pytorch_jacinto_ai.engine import train_pixel2pixel
  64
  65 # Create the parser and set default arguments
  66 args = train_pixel2pixel.get_config()
  67
  68 ################################
  69 #Modify arguments
  70 args.model_name = 'deeplabv3lite_mobilenetv2_tv' #'deeplabv3lite_mobilenetv2_tv' #'fpn_pixel2pixel_aspp_mobilenetv2_tv' #'fpn_pixel2pixel_aspp_resnet50'
  71 args.dataset_name = 'cityscapes_segmentation' #'cityscapes_segmentation' #'voc_segmentation'
  72
  73 args.data_path = './data/datasets/cityscapes/data' #'./data/datasets/cityscapes/data' #'./data/datasets/voc'
  74
  75 #args.save_path = './data/checkpoints'
  76
  77 args.pretrained = 'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth'
  78                         # 'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth'
  79                         # './data/modelzoo/pretrained/pytorch/imagenet_classification/ericsun99/MobileNet-V2-Pytorch/mobilenetv2_Top1_71.806_Top2_90.410.pth.tar'
  80                         # 'https://download.pytorch.org/models/resnet50-19c8e357.pth'
  81
  82 #args.resume = './data/checkpoints/cityscapes_segmentation/2019-04-11-05-35-55_cityscapes_segmentation_deeplabv3lite_mobilenetv2_relu_resize768x384_traincrop768x384/checkpoint.pth.tar'
  83
  84 args.model_config.input_channels = (3,)
  85 args.model_config.output_type = ['segmentation']
  86 args.model_config.output_channels = None
  87 args.model_config.output_range = None
  88 args.model_config.num_decoders = None                       #0, 1, None
  89
  90 args.losses = [['segmentation_loss']]
  91 args.metrics = [['segmentation_metrics']]
  92
  93 args.solver = 'adam'                                        #'sgd' #'adam'
  94 args.epochs = 250                                           #200
  95 args.epoch_size = 0                                         #0 #0.5
  96 args.epoch_size_val = 0                                     #0 #0.5
  97 args.scheduler = 'step'                                     #'poly' #'step'
  98 args.multistep_gamma = 0.25                                 #0.5 #only for step scheduler
  99 args.milestones = (100, 200)                                #only for step scheduler
 100 args.polystep_power = 0.9                                   #only for poly scheduler
 101 args.iter_size = 1                                          #2
 102
 103 args.lr = 4e-4                                              #1e-4 #0.01 #7e-3 #1e-4 #2e-4
 104 args.batch_size = 12                                        #12 #16 #32 #64
 105 args.weight_decay = 1e-4                                    #1e-4  #4e-5 #1e-5
 106
 107 args.img_resize = (384, 768)                                #(384, 768) (512, 1024) #(1024, 2048)
 108 args.output_size = (1024, 2048)                             #target output size for evaluation
 109
 110 args.transform_rotation = 5                                 #rotation degrees
 111
 112 #args.image_mean = [123.675, 116.28, 103.53]
 113 #args.image_scale = [0.017125, 0.017507, 0.017429]
 114
 115 #args.parallel_model=False
 116 #args.print_model = True
 117 #args.generate_onnx = False
 118 #args.run_soon = False
 119 #args.evaluate_start = False
 120 args.print_freq = 10
 121
 122 #args.phase = 'validation' #'training'
 123 #args.quantize = True
 124 #args.per_channel_q = True
 125
 126
 127 # defining date from outside can help to write multiple phases into the same folder
 128 args.date = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
 129
 130 ################################
 131 # set other args
 132 for key in vars(cmds):
 133     if key == 'gpus':
 134         pass # already taken care above, since this has to be done before importing pytorch
 135     elif hasattr(args, key):
 136         value = getattr(cmds, key)
 137         if value != 'None' and value is not None:
 138             setattr(args, key, value)
 139     else:
 140         assert False, f'invalid argument {key}'
 141 #
 142
 143 ################################
 144 # Run the given phase
 145 train_pixel2pixel.main(args)
 146
 147 ################################
 148 # In addition run a quantization aware training, starting from the trained model
 149 if 'training' in args.phase and (not args.quantize):
 150     save_path = train_pixel2pixel.get_save_path(args)
 151     args.pretrained = os.path.join(save_path, 'model_best.pth.tar') if (args.epochs>0) else args.pretrained
 152     args.phase = 'training_quantize'
 153     args.quantize = True
 154     args.lr = 1e-5
 155     args.epochs = 50
 156     # quantized training will use only one GPU in the engine - so reduce the batch_size
 157     num_gpus = len(str(os.environ["CUDA_VISIBLE_DEVICES"]).split(',')) if ("CUDA_VISIBLE_DEVICES" in os.environ) else None
 158     args.batch_size = (args.batch_size//num_gpus) if (num_gpus is not None) else args.batch_size
 159     train_pixel2pixel.main(args)
 160 #
 161
 162 ################################
 163 # In addition run a separate validation
 164 if 'training' in args.phase or 'calibration' in args.phase:
 165     save_path = train_pixel2pixel.get_save_path(args)
 166     args.pretrained = os.path.join(save_path, 'model_best.pth.tar')
 167     args.phase = 'validation'
 168     args.quantize = True
 169     train_pixel2pixel.main(args)
 170 #