9686aa3d1f7eb7f8272aab74c1b0e75658a41bd6
[jacinto-ai/pytorch-jacinto-ai-devkit.git] / scripts / train_motion_segmentation_main.py
1 import sys
2 import os
3 import cv2
4 import argparse
5 import datetime
6 import numpy as np
8 ################################
9 from pytorch_jacinto_ai.xnn.utils import str2bool, splitstr2bool
10 parser = argparse.ArgumentParser()
11 parser.add_argument('--save_path', type=str, default=None, help='checkpoint save folder')
12 parser.add_argument('--gpus', type=int, nargs='*', default=None, help='Base learning rate')
13 parser.add_argument('--batch_size', type=int, default=None, help='Batch size')
14 parser.add_argument('--lr', type=float, default=None, help='Base learning rate')
15 parser.add_argument('--lr_clips', type=float, default=None, help='Learning rate for clips in PAct2')
16 parser.add_argument('--lr_calib', type=float, default=None, help='Learning rate for calibration')
17 parser.add_argument('--model_name', type=str, default=None, help='model name')
18 parser.add_argument('--dataset_name', type=str, default=None, help='dataset name')
19 parser.add_argument('--image_folders', type=str, default=('leftImg8bit_flow_confidence', 'leftImg8bit'), nargs='*', help='image_folders')
20 parser.add_argument('--is_flow', type=splitstr2bool, default=[[False, False],[False]] ,nargs='*', help='whether any input or target is optical flow')
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()
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 # to avoid hangs in data loader with multi threads
55 # this was observed after using cv2 image processing functions
56 # https://github.com/pytorch/pytorch/issues/1355
57 cv2.setNumThreads(0)
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
65 #Create the parse and set default arguments
66 args = train_pixel2pixel.get_config()
69 #Modify arguments
71 args.model_name = 'deeplabv3lite_mobilenetv2_tv'
73 #args.save_path = './data/checkpoints'
74 args.dataset_name = 'cityscapes_motion_multi_input'
75 args.model_config.output_type = ['segmentation']
76 args.model_config.output_channels = None
77 args.losses = [['segmentation_loss']]
78 args.metrics = [['segmentation_metrics']]
80 args.data_path = './data/datasets/cityscapes_768x384/data'  #'./data/datasets/cityscapes/data'  #'./data/tiad/data'
82 args.pretrained = "/user/a0132471/Files/pytorch/pytorch-devkit/data/modelzoo/semantic_segmentation/cityscapes/deeplabv3lite-mobilenetv2/cityscapes_segmentation_deeplabv3lite-mobilenetv2_2019-06-26-08-59-32.pth"
83                         #'./data/modelzoo/pretrained/pytorch/cityscapes_segmentation/v0.9-2018-12-07-19:38:26_cityscapes_segmentation_deeplabv3lite_mobilenetv2_relu_resize768x384_traincrop768x384_(68.9%)/model_best.pth.tar'
84                         #'./data/checkpoints/cityscapes_segmentation/2019-01-18-15:41:05_cityscapes_segmentation_deeplabv3lite_mobilenetv2_x2p0_relu_resize768x384_traincrop768x384/model_best.pth.tar'
85                         #'./data/checkpoints/Models for Deepak/cityscapes_motion_image_rawdof_rgb/2018-11-13-17:54:06_cityscapes_motion_image_rawdof_rgb_deeplabv3lite_mobilenetv2_dualstream_resize1024x512_traincrop512x512_mIOU_86.4/model_best.pth.tar'
88 args.solver = 'adam'                    #'sgd' #'adam'
89 args.epochs = 250                       #300
90 args.epoch_size = 0                     #0 #0.5
91 args.scheduler = 'step'                 #'poly' #'step'
92 args.multistep_gamma = 0.5              #only for step scheduler
93 args.milestones = [100,  200]           #only for step scheduler
94 args.polystep_power = 0.9               #only for poly scheduler
95 args.iter_size = 1                      #2
96 #args.phase = 'validation'
98 args.lr = 1e-4                          #1e-4 #0.01 #7e-3 #1e-4 #2e-4
99 args.batch_size = 16                    #12 #16 #32 #64
100 args.weight_decay = 1e-4                #4e-5 #1e-5
102 args.img_resize = (384, 768)            #(512, 1024) #(1024, 2048)
103 args.rand_scale = (1.0, 2.0)            #(1.0,2.0)
104 args.rand_crop = (384, 768)             #(512,512) #(512,1024)
105 args.output_size = (1024, 2048)         #target output size for evaluation
107 args.transform_rotation = 5             #0  #rotation degrees
109 args.model_config.input_channels = (3,3)
110 args.model_config.aspp_dil = (2, 4, 6)
112 # TODO: this is not clean - fix it later
113 args.dataset_config.image_folders = cmds.image_folders
115 args.generate_onnx = True
116 #args.evaluate = True
117 #args.quantize = True
118 #args.generate_onnx = True
121 ################################
122 for key in vars(cmds):
123     if (key == 'gpus') | (key == 'image_folders'):
124         pass # already taken care above, since this has to be done before importing pytorch
125     elif hasattr(args, key):
126         value = getattr(cmds, key)
127         if value != 'None' and value is not None:
128             setattr(args, key, value)
129     elif hasattr(args.model_config, key):
130         value = getattr(cmds, key)
131         if value != 'None' and value is not None:
132             setattr(args.model_config, key, value)
133     elif hasattr(args.dataset_config, key):
134         value = getattr(cmds, key)
135         if value != 'None' and value is not None:
136             setattr(args.dataset_config, key, value)
137     else:
138         assert False, f'invalid argument {key}'
141 ################################
142 # Run the given phase
143 train_pixel2pixel.main(args)
145 ################################
146 # In addition run a quantization aware training, starting from the trained model
147 if 'training' in args.phase and (not args.quantize):
148     save_path = train_pixel2pixel.get_save_path(args)
149     args.pretrained = os.path.join(save_path, 'model_best.pth.tar') if (args.epochs>0) else args.pretrained
150     args.phase = 'training_quantize'
151     args.quantize = True
152     args.lr = 1e-5
153     args.epochs = 25
154     # quantized training will use only one GPU in the engine - so reduce the batch_size
155     num_gpus = len(str(os.environ["CUDA_VISIBLE_DEVICES"]).split(',')) if ("CUDA_VISIBLE_DEVICES" in os.environ) else None
156     args.batch_size = (args.batch_size//num_gpus) if (num_gpus is not None) else args.batch_size
157     train_pixel2pixel.main(args)
160 ################################
161 # In addition run a separate validation
162 if 'training' in args.phase or 'calibration' in args.phase:
163     save_path = train_pixel2pixel.get_save_path(args)
164     args.pretrained = os.path.join(save_path, 'model_best.pth.tar')
165     args.phase = 'validation'
166     args.quantize = True
167     train_pixel2pixel.main(args)