def dataloader(bs, gpus): roidb, ratio_list, ratio_index = \ combined_roidb_for_training_semseg('cityscapes_semseg_val') sampler = MinibatchSampler(ratio_list, ratio_index) dataset = RoiDataLoader(roidb, 19, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_size=bs, sampler=sampler, num_workers=gpus, collate_fn=collate_minibatch_semseg) return dataloader
def dataloader(bs, gpus): roidb, ratio_list, ratio_index = \ combined_roidb_for_training_semseg('cityscapes_semseg_train') sampler = MinibatchSampler(ratio_list, ratio_index) dataset = RoiDataLoader(roidb, 19, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_size=bs, sampler=sampler, num_workers=gpus, collate_fn=collate_minibatch_semseg) return dataloader return torch.randn(bs*gpus, 3, 720, 720), \ torch.LongTensor(np.random.randint(0, 19, (bs*gpus, 90, 90), dtype=np.long))
def main(): """Main function""" args = parse_args() print('Called with args:') print(args) if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") if args.cuda or cfg.NUM_GPUS > 0: cfg.CUDA = True else: raise ValueError("Need Cuda device to run !") if args.dataset == "coco2017": cfg.TRAIN.DATASETS = ('coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == "keypoints_coco2017": cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "common": cfg.TRAIN.DATASETS = ('common_train', ) cfg.MODEL.NUM_CLASSES = 81 else: raise ValueError("Unexpected args.dataset: {}".format(args.dataset)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) ### Adaptively adjust some configs ### original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH if args.batch_size is None: args.batch_size = original_batch_size cfg.NUM_GPUS = torch.cuda.device_count() assert (args.batch_size % cfg.NUM_GPUS) == 0, \ 'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS) cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS print('Batch size change from {} (in config file) to {}'.format( original_batch_size, args.batch_size)) print('NUM_GPUs: %d, TRAIN.IMS_PER_BATCH: %d' % (cfg.NUM_GPUS, cfg.TRAIN.IMS_PER_BATCH)) if args.num_workers is not None: cfg.DATA_LOADER.NUM_THREADS = args.num_workers print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS) ### Adjust learning based on batch size change linearly old_base_lr = cfg.SOLVER.BASE_LR cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size print('Adjust BASE_LR linearly according to batch size change: {} --> {}'. format(old_base_lr, cfg.SOLVER.BASE_LR)) ### Overwrite some solver settings from command line arguments if args.optimizer is not None: cfg.SOLVER.TYPE = args.optimizer if args.lr is not None: cfg.SOLVER.BASE_LR = args.lr if args.lr_decay_gamma is not None: cfg.SOLVER.GAMMA = args.lr_decay_gamma timers = defaultdict(Timer) ### Dataset ### timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() train_size = len(roidb) logger.info('{:d} roidb entries'.format(train_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) sampler = MinibatchSampler(ratio_list, ratio_index) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, sampler=sampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) assert_and_infer_cfg() ### Model ### maskRCNN = Generalized_RCNN() if cfg.CUDA: maskRCNN.cuda() ### Optimizer ### bias_params = [] nonbias_params = [] for key, value in dict(maskRCNN.named_parameters()).items(): if value.requires_grad: if 'bias' in key: bias_params.append(value) else: nonbias_params.append(value) params = [{ 'params': nonbias_params, 'lr': cfg.SOLVER.BASE_LR, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': bias_params, 'lr': cfg.SOLVER.BASE_LR * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }] if cfg.SOLVER.TYPE == "SGD": optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizer = torch.optim.Adam(params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) if args.resume: assert checkpoint['iters_per_epoch'] == train_size // args.batch_size, \ "iters_per_epoch should match for resume" # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. # optimizer.load_state_dict(checkpoint['optimizer']) misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) if checkpoint['step'] == (checkpoint['iters_per_epoch'] - 1): # Resume from end of an epoch args.start_epoch = checkpoint['epoch'] + 1 args.start_iter = 0 else: # Resume from the middle of an epoch. # NOTE: dataloader is not synced with previous state args.start_epoch = checkpoint['epoch'] args.start_iter = checkpoint['step'] + 1 del checkpoint torch.cuda.empty_cache() if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values) logging.info("loading Detectron weights %s", args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) lr = optimizer.param_groups[0][ 'lr'] # lr of non-bias parameters, for commmand line outputs. maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True) ### Training Setups ### args.run_name = misc_utils.get_run_name() output_dir = misc_utils.get_output_dir(args, args.run_name) args.cfg_filename = os.path.basename(args.cfg_file) if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) ### Training Loop ### maskRCNN.train() training_stats = TrainingStats( args, args.disp_interval, tblogger if args.use_tfboard and not args.no_save else None) iters_per_epoch = int(train_size / args.batch_size) # drop last args.iters_per_epoch = iters_per_epoch ckpt_interval_per_epoch = iters_per_epoch // args.ckpt_num_per_epoch try: logger.info('Training starts !') args.step = args.start_iter global_step = iters_per_epoch * args.start_epoch + args.step for args.epoch in range(args.start_epoch, args.start_epoch + args.num_epochs): # ---- Start of epoch ---- # adjust learning rate if args.lr_decay_epochs and args.epoch == args.lr_decay_epochs[ 0] and args.start_iter == 0: args.lr_decay_epochs.pop(0) net_utils.decay_learning_rate(optimizer, lr, cfg.SOLVER.GAMMA) lr *= cfg.SOLVER.GAMMA for args.step, input_data in zip( range(args.start_iter, iters_per_epoch), dataloader): for key in input_data: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list(map(Variable, input_data[key])) training_stats.IterTic() net_outputs = maskRCNN(**input_data) training_stats.UpdateIterStats(net_outputs) loss = net_outputs['total_loss'] optimizer.zero_grad() loss.backward() optimizer.step() training_stats.IterToc() if (args.step + 1) % ckpt_interval_per_epoch == 0: net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer) if args.step % args.disp_interval == 0: log_training_stats(training_stats, global_step, lr) global_step += 1 # ---- End of epoch ---- # save checkpoint net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer) # reset starting iter number after first epoch args.start_iter = 0 # ---- Training ends ---- if iters_per_epoch % args.disp_interval != 0: # log last stats at the end log_training_stats(training_stats, global_step, lr) except (RuntimeError, KeyboardInterrupt): logger.info('Save ckpt on exception ...') net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer) logger.info('Save ckpt done.') stack_trace = traceback.format_exc() print(stack_trace) finally: if args.use_tfboard and not args.no_save: tblogger.close()
def main(): """Main function""" args = parse_args() print('Called with args:') print(args) if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") if args.cuda or cfg.NUM_GPUS > 0: cfg.CUDA = True else: raise ValueError("Need Cuda device to run !") if args.dataset == "coco2017": cfg.TRAIN.DATASETS = ('coco_2017_train',) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == "keypoints_coco2017": cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train',) cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError("Unexpected args.dataset: {}".format(args.dataset)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) ### Adaptively adjust some configs ### original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH original_num_gpus = cfg.NUM_GPUS if args.batch_size is None: args.batch_size = original_batch_size cfg.NUM_GPUS = torch.cuda.device_count() assert (args.batch_size % cfg.NUM_GPUS) == 0, \ 'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS) cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS effective_batch_size = args.iter_size * args.batch_size print('effective_batch_size = batch_size * iter_size = %d * %d' % (args.batch_size, args.iter_size)) print('Adaptive config changes:') print(' effective_batch_size: %d --> %d' % (original_batch_size, effective_batch_size)) print(' NUM_GPUS: %d --> %d' % (original_num_gpus, cfg.NUM_GPUS)) print(' IMS_PER_BATCH: %d --> %d' % (original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH)) ### Adjust learning based on batch size change linearly # For iter_size > 1, gradients are `accumulated`, so lr is scaled based # on batch_size instead of effective_batch_size old_base_lr = cfg.SOLVER.BASE_LR cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size print('Adjust BASE_LR linearly according to batch_size change:\n' ' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR)) ### Adjust solver steps step_scale = original_batch_size / effective_batch_size old_solver_steps = cfg.SOLVER.STEPS old_max_iter = cfg.SOLVER.MAX_ITER cfg.SOLVER.STEPS = list(map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS)) cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5) print('Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n' ' SOLVER.STEPS: {} --> {}\n' ' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps, cfg.SOLVER.STEPS, old_max_iter, cfg.SOLVER.MAX_ITER)) # Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function # of `collect_and_distribute_fpn_rpn_proposals.py` # # post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5) if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN: cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch print('Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n' ' cfg.FPN.RPN_COLLECT_SCALE: {}'.format(cfg.FPN.RPN_COLLECT_SCALE)) if args.num_workers is not None: cfg.DATA_LOADER.NUM_THREADS = args.num_workers print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS) ### Overwrite some solver settings from command line arguments if args.optimizer is not None: cfg.SOLVER.TYPE = args.optimizer if args.lr is not None: cfg.SOLVER.BASE_LR = args.lr if args.lr_decay_gamma is not None: cfg.SOLVER.GAMMA = args.lr_decay_gamma assert_and_infer_cfg() timers = defaultdict(Timer) ### Dataset ### timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) # Effective training sample size for one epoch train_size = roidb_size // args.batch_size * args.batch_size batchSampler = BatchSampler( sampler=MinibatchSampler(ratio_list, ratio_index), batch_size=args.batch_size, drop_last=True ) dataset = RoiDataLoader( roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_sampler=batchSampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) dataiterator = iter(dataloader) ### Model ### maskRCNN = Generalized_RCNN() if cfg.CUDA: maskRCNN.cuda() ### Optimizer ### gn_param_nameset = set() for name, module in maskRCNN.named_modules(): if isinstance(module, nn.GroupNorm): gn_param_nameset.add(name+'.weight') gn_param_nameset.add(name+'.bias') gn_params = [] gn_param_names = [] bias_params = [] bias_param_names = [] nonbias_params = [] nonbias_param_names = [] nograd_param_names = [] for key, value in maskRCNN.named_parameters(): if value.requires_grad: if 'bias' in key: bias_params.append(value) bias_param_names.append(key) elif key in gn_param_nameset: gn_params.append(value) gn_param_names.append(key) else: nonbias_params.append(value) nonbias_param_names.append(key) else: nograd_param_names.append(key) assert (gn_param_nameset - set(nograd_param_names) - set(bias_param_names)) == set(gn_param_names) # Learning rate of 0 is a dummy value to be set properly at the start of training params = [ {'params': nonbias_params, 'lr': 0, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY}, {'params': bias_params, 'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0}, {'params': gn_params, 'lr': 0, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY_GN} ] # names of paramerters for each paramter param_names = [nonbias_param_names, bias_param_names, gn_param_names] if cfg.SOLVER.TYPE == "SGD": optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizer = torch.optim.Adam(params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) if args.resume: args.start_step = checkpoint['step'] + 1 if 'train_size' in checkpoint: # For backward compatibility if checkpoint['train_size'] != train_size: print('train_size value: %d different from the one in checkpoint: %d' % (train_size, checkpoint['train_size'])) # reorder the params in optimizer checkpoint's params_groups if needed # misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint) # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. optimizer.load_state_dict(checkpoint['optimizer']) # misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) del checkpoint torch.cuda.empty_cache() if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values) logging.info("loading Detectron weights %s", args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) lr = optimizer.param_groups[0]['lr'] # lr of non-bias parameters, for commmand line outputs. maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True) ### Training Setups ### args.run_name = misc_utils.get_run_name() + '_step' output_dir = misc_utils.get_output_dir(args, args.run_name) args.cfg_filename = os.path.basename(args.cfg_file) if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) ### Training Loop ### maskRCNN.train() CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS) # Set index for decay steps decay_steps_ind = None for i in range(1, len(cfg.SOLVER.STEPS)): if cfg.SOLVER.STEPS[i] >= args.start_step: decay_steps_ind = i break if decay_steps_ind is None: decay_steps_ind = len(cfg.SOLVER.STEPS) training_stats = TrainingStats( args, args.disp_interval, tblogger if args.use_tfboard and not args.no_save else None) try: logger.info('Training starts !') step = args.start_step for step in range(args.start_step, cfg.SOLVER.MAX_ITER): # Warm up if step < cfg.SOLVER.WARM_UP_ITERS: method = cfg.SOLVER.WARM_UP_METHOD if method == 'constant': warmup_factor = cfg.SOLVER.WARM_UP_FACTOR elif method == 'linear': alpha = step / cfg.SOLVER.WARM_UP_ITERS warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 - alpha) + alpha else: raise KeyError('Unknown SOLVER.WARM_UP_METHOD: {}'.format(method)) lr_new = cfg.SOLVER.BASE_LR * warmup_factor net_utils.update_learning_rate(optimizer, lr, lr_new) lr = optimizer.param_groups[0]['lr'] assert lr == lr_new elif step == cfg.SOLVER.WARM_UP_ITERS: net_utils.update_learning_rate(optimizer, lr, cfg.SOLVER.BASE_LR) lr = optimizer.param_groups[0]['lr'] assert lr == cfg.SOLVER.BASE_LR # Learning rate decay if decay_steps_ind < len(cfg.SOLVER.STEPS) and \ step == cfg.SOLVER.STEPS[decay_steps_ind]: logger.info('Decay the learning on step %d', step) lr_new = lr * cfg.SOLVER.GAMMA net_utils.update_learning_rate(optimizer, lr, lr_new) lr = optimizer.param_groups[0]['lr'] assert lr == lr_new decay_steps_ind += 1 training_stats.IterTic() optimizer.zero_grad() for inner_iter in range(args.iter_size): try: input_data = next(dataiterator) except StopIteration: dataiterator = iter(dataloader) input_data = next(dataiterator) for key in input_data: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list(map(Variable, input_data[key])) net_outputs = maskRCNN(**input_data) training_stats.UpdateIterStats(net_outputs, inner_iter) loss = net_outputs['total_loss'] loss.backward() optimizer.step() training_stats.IterToc() training_stats.LogIterStats(step, lr) if (step+1) % CHECKPOINT_PERIOD == 0: save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) # ---- Training ends ---- # Save last checkpoint save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) except (RuntimeError, KeyboardInterrupt): del dataiterator logger.info('Save ckpt on exception ...') save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) logger.info('Save ckpt done.') stack_trace = traceback.format_exc() print(stack_trace) finally: if args.use_tfboard and not args.no_save: tblogger.close()
def main(): """Main function""" args = parse_args() print('Called with args:') print(args) if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") if args.cuda or cfg.NUM_GPUS > 0: cfg.CUDA = True else: raise ValueError("Need Cuda device to run !") if args.dataset == "coco2017": cfg.TRAIN.DATASETS = ('coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == "keypoints_coco2017": cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError("Unexpected args.dataset: {}".format(args.dataset)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) ### Adaptively adjust some configs ### original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH if args.batch_size is None: args.batch_size = original_batch_size cfg.NUM_GPUS = torch.cuda.device_count() assert (args.batch_size % cfg.NUM_GPUS) == 0, \ 'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS) cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS print('Batch size change from {} (in config file) to {}'.format( original_batch_size, args.batch_size)) print('NUM_GPUs: %d, TRAIN.IMS_PER_BATCH: %d' % (cfg.NUM_GPUS, cfg.TRAIN.IMS_PER_BATCH)) if args.num_workers is not None: cfg.DATA_LOADER.NUM_THREADS = args.num_workers print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS) ### Adjust learning based on batch size change linearly old_base_lr = cfg.SOLVER.BASE_LR cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size print('Adjust BASE_LR linearly according to batch size change: {} --> {}'. format(old_base_lr, cfg.SOLVER.BASE_LR)) ### Overwrite some solver settings from command line arguments if args.optimizer is not None: cfg.SOLVER.TYPE = args.optimizer if args.lr is not None: cfg.SOLVER.BASE_LR = args.lr if args.lr_decay_gamma is not None: cfg.SOLVER.GAMMA = args.lr_decay_gamma assert_and_infer_cfg() timers = defaultdict(Timer) ### Dataset ### timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) # Effective training sample size for one epoch train_size = roidb_size // args.batch_size * args.batch_size sampler = MinibatchSampler(ratio_list, ratio_index) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, drop_last=True, sampler=sampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) dataiterator = iter(dataloader) ### Model ### maskRCNN = Generalized_RCNN() if cfg.CUDA: maskRCNN.cuda() ### Optimizer ### bias_params = [] nonbias_params = [] for key, value in dict(maskRCNN.named_parameters()).items(): if value.requires_grad: if 'bias' in key: bias_params.append(value) else: nonbias_params.append(value) # Learning rate of 0 is a dummy value to be set properly at the start of training params = [{ 'params': nonbias_params, 'lr': 0, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': bias_params, 'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }] if cfg.SOLVER.TYPE == "SGD": optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizer = torch.optim.Adam(params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) if args.resume: args.start_step = checkpoint['step'] + 1 assert checkpoint['train_size'] == train_size # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. # optimizer.load_state_dict(checkpoint['optimizer']) misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) del checkpoint torch.cuda.empty_cache() if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values) logging.info("loading Detectron weights %s", args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) lr = optimizer.param_groups[0][ 'lr'] # lr of non-bias parameters, for commmand line outputs. maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True) ### Training Setups ### run_name = misc_utils.get_run_name() output_dir = misc_utils.get_output_dir(args, run_name) if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) ### Training Loop ### maskRCNN.train() CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS) # Set index for decay steps decay_steps_ind = None for i in range(1, len(cfg.SOLVER.STEPS)): if cfg.SOLVER.STEPS[i] > args.start_step: decay_steps_ind = i if decay_steps_ind is None: decay_steps_ind = len(cfg.SOLVER.STEPS) logger.info('Training starts !') loss_avg = 0 try: timers['train_loop'].tic() for step in range(args.start_step, cfg.SOLVER.MAX_ITER): # Warm up if step < cfg.SOLVER.WARM_UP_ITERS: method = cfg.SOLVER.WARM_UP_METHOD if method == 'constant': warmup_factor = cfg.SOLVER.WARM_UP_FACTOR elif method == 'linear': alpha = step / cfg.SOLVER.WARM_UP_ITERS warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 - alpha) + alpha else: raise KeyError( 'Unknown SOLVER.WARM_UP_METHOD: {}'.format(method)) lr_new = cfg.SOLVER.BASE_LR * warmup_factor net_utils.update_learning_rate(optimizer, lr, lr_new) lr = lr_new elif step == cfg.SOLVER.WARM_UP_ITERS: lr = cfg.SOLVER.BASE_LR # Learning rate decay if decay_steps_ind < len(cfg.SOLVER.STEPS) and \ step == cfg.SOLVER.STEPS[decay_steps_ind]: logger.info('Decay the learning on step %d', step) lr_new = lr * cfg.SOLVER.GAMMA net_utils.update_learning_rate(optimizer, lr, lr_new) lr = lr_new decay_steps_ind += 1 try: input_data = next(dataiterator) except StopIteration: dataiterator = iter(dataloader) input_data = next(dataiterator) for key in input_data: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list(map(Variable, input_data[key])) outputs = maskRCNN(**input_data) rois_label = outputs['rois_label'] cls_score = outputs['cls_score'] bbox_pred = outputs['bbox_pred'] loss_rpn_cls = outputs['loss_rpn_cls'].mean() loss_rpn_bbox = outputs['loss_rpn_bbox'].mean() loss_rcnn_cls = outputs['loss_rcnn_cls'].mean() loss_rcnn_bbox = outputs['loss_rcnn_bbox'].mean() loss = loss_rpn_cls + loss_rpn_bbox + loss_rcnn_cls + loss_rcnn_bbox if cfg.MODEL.MASK_ON: loss_rcnn_mask = outputs['loss_rcnn_mask'].mean() loss += loss_rcnn_mask if cfg.MODEL.KEYPOINTS_ON: loss_rcnn_keypoints = outputs['loss_rcnn_keypoints'].mean() loss += loss_rcnn_keypoints loss_avg += loss.data.cpu().numpy()[0] optimizer.zero_grad() loss.backward() optimizer.step() if (step + 1) % CHECKPOINT_PERIOD == 0: save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) if ((step % args.disp_interval == 0 and (step - args.start_step >= args.disp_interval)) or step == cfg.SOLVER.MAX_ITER - 1): diff = timers['train_loop'].toc(average=False) loss_avg /= args.disp_interval loss_rpn_cls = loss_rpn_cls.data[0] loss_rpn_bbox = loss_rpn_bbox.data[0] loss_rcnn_cls = loss_rcnn_cls.data[0] loss_rcnn_bbox = loss_rcnn_bbox.data[0] fg_cnt = torch.sum(rois_label.data.ne(0)) bg_cnt = rois_label.data.numel() - fg_cnt print("[ %s ][ step %d ]" % (run_name, step)) print("\t\tloss: %.4f, lr: %.2e" % (loss_avg, lr)) print("\t\tfg/bg=(%d/%d), time cost: %f" % (fg_cnt, bg_cnt, diff)) print( "\t\trpn_cls: %.4f, rpn_bbox: %.4f, rcnn_cls: %.4f, rcnn_bbox %.4f" % (loss_rpn_cls, loss_rpn_bbox, loss_rcnn_cls, loss_rcnn_bbox)) print_prefix = "\t\t" if cfg.MODEL.MASK_ON: loss_rcnn_mask = loss_rcnn_mask.data[0] print("%srcnn_mask %.4f" % (print_prefix, loss_rcnn_mask)) print_prefix = ", " if cfg.MODEL.KEYPOINTS_ON: loss_rcnn_keypoints = loss_rcnn_keypoints.data[0] print("%srcnn_keypoints %.4f" % (print_prefix, loss_rcnn_keypoints)) if args.use_tfboard and not args.no_save: info = { 'lr': lr, 'loss': loss_avg, 'loss_rpn_cls': loss_rpn_cls, 'loss_rpn_box': loss_rpn_bbox, 'loss_rcnn_cls': loss_rcnn_cls, 'loss_rcnn_box': loss_rcnn_bbox, } if cfg.MODEL.MASK_ON: info['loss_rcnn_mask'] = loss_rcnn_mask if cfg.MODEL.KEYPOINTS_ON: info['loss_rcnn_keypoints'] = loss_rcnn_keypoints for tag, value in info.items(): tblogger.add_scalar(tag, value, step) loss_avg = 0 timers['train_loop'].tic() # Save last checkpoint save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) except (RuntimeError, KeyboardInterrupt) as e: print('Save on exception:', e) save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) stack_trace = traceback.format_exc() print(stack_trace) finally: # ---- Training ends ---- if args.use_tfboard and not args.no_save: tblogger.close()
from roi_data.loader import RoiDataLoader, MinibatchSampler import torch import numpy as np ratio_list = np.array([1, 2, 3, 4, 5]) ratio_index = np.array([0, 1, 2, 3, 4]) sampler = MinibatchSampler(ratio_list, ratio_index) batch_sampler = iter( torch.utils.data.sampler.BatchSampler(sampler, batch_size=2, drop_last=False)) for i in range(10): items = [] for ids, _ in next(batch_sampler): items.append(ids) print(i, items)
def main(): saveNetStructure = False """Main function""" args = parse_args() print('Called with args:') print(args) if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") if args.cuda or cfg.NUM_GPUS > 0: #set gpu device os.environ["CUDA_VISIBLE_DEVICES"] = ",".join( [str(ids) for ids in args.device_ids]) torch.backends.cudnn.benchmark = True cfg.CUDA = True else: raise ValueError("Need Cuda device to run !") if args.dataset == "coco2017": cfg.TRAIN.DATASETS = ('coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == "keypoints_coco2017": cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cityscapes": cfg.TRAIN.DATASETS = ('cityscapes_semseg_train', ) cfg.MODEL.NUM_CLASSES = 19 else: raise ValueError("Unexpected args.dataset: {}".format(args.dataset)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) ### Adaptively adjust some configs ### original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH if args.batch_size is None: args.batch_size = original_batch_size cfg.NUM_GPUS = torch.cuda.device_count() assert (args.batch_size % cfg.NUM_GPUS) == 0, \ 'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS) cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS print('Batch size change from {} (in config file) to {}'.format( original_batch_size, args.batch_size)) print('NUM_GPUs: %d, TRAIN.IMS_PER_BATCH: %d' % (cfg.NUM_GPUS, cfg.TRAIN.IMS_PER_BATCH)) if args.num_workers is not None: cfg.DATA_LOADER.NUM_THREADS = args.num_workers print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS) ### Adjust learning based on batch size change linearly old_base_lr = cfg.SOLVER.BASE_LR cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size print('Adjust BASE_LR linearly according to batch size change: {} --> {}'. format(old_base_lr, cfg.SOLVER.BASE_LR)) ### Overwrite some solver settings from command line arguments if args.optimizer is not None: cfg.SOLVER.TYPE = args.optimizer if args.lr is not None: cfg.SOLVER.BASE_LR = args.lr if args.lr_decay_gamma is not None: cfg.SOLVER.GAMMA = args.lr_decay_gamma timers = defaultdict(Timer) ### Dataset ### timers['roidb'].tic() if cfg.SEM.SEM_ON or cfg.DISP.DISP_ON: roidb, ratio_list, ratio_index = combined_roidb_for_training_semseg( cfg.TRAIN.DATASETS) else: roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() train_size = len(roidb) logger.info('{:d} roidb entries'.format(train_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) sampler = MinibatchSampler(ratio_list, ratio_index) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, sampler=sampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch_semseg if cfg.SEM.SEM_ON or cfg.DISP.DISP_ON else collate_minibatch) assert_and_infer_cfg() #for args.step, input_data in zip(range(100), dataloader): # data_L = input_data['data'] # data_R = input_data['data_R'] # label = input_data['disp_label_0'] # cv2.imwrite('ims_L.png', data_L[0].numpy()[0].transpose(1,2,0)[:,:,::-1]+cfg.PIXEL_MEANS) # cv2.imwrite('ims_R.png', data_R[0].numpy()[0].transpose(1,2,0)[:,:,::-1]+cfg.PIXEL_MEANS) # cv2.imwrite('label.png', label[0].numpy()[0]) # return ### Model ### dispSeg = DispSeg() if cfg.CUDA: dispSeg.to('cuda') pspnet_bias_params = [] pspnet_nonbias_params = [] for key, value in dict(dispSeg.pspnet.named_parameters()).items(): if value.requires_grad: if 'bias' in key: pspnet_bias_params.append(value) else: pspnet_nonbias_params.append(value) pspnet_params = [{ 'params': pspnet_nonbias_params, 'lr': cfg.SOLVER.BASE_LR, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': pspnet_bias_params, 'lr': cfg.SOLVER.BASE_LR * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }] glassGCN_bias_params = [] glassGCN_nonbias_params = [] for key, value in dict(dispSeg.glassGCN.named_parameters()).items(): if value.requires_grad: if 'bias' in key: glassGCN_bias_params.append(value) else: glassGCN_nonbias_params.append(value) segdisp3d_params = [{ 'params': glassGCN_nonbias_params, 'lr': cfg.SOLVER.BASE_LR, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': glassGCN_bias_params, 'lr': cfg.SOLVER.BASE_LR * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }] if cfg.SOLVER.TYPE == "SGD": optimizerP = torch.optim.SGD(pspnet_params, momentum=cfg.SOLVER.MOMENTUM) optimizerS = torch.optim.SGD(segdisp3d_params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizerP = torch.optim.Adam(pspnet_params) optimizerS = torch.optim.Adam(segdisp3d_params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(pspnet, checkpoint['model']) net_utils.load_ckpt(segdisp3d, checkpoint['model']) if args.resume: assert checkpoint['iters_per_epoch'] == train_size // args.batch_size, \ "iters_per_epoch should match for resume" # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. # optimizer.load_state_dict(checkpoint['optimizer']) misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) if checkpoint['step'] == (checkpoint['iters_per_epoch'] - 1): # Resume from end of an epoch args.start_epoch = checkpoint['epoch'] + 1 args.start_iter = 0 else: # Resume from the middle of an epoch. # NOTE: dataloader is not synced with previous state args.start_epoch = checkpoint['epoch'] args.start_iter = checkpoint['step'] + 1 del checkpoint torch.cuda.empty_cache() lr = optimizerP.param_groups[0][ 'lr'] # lr of non-bias parameters, for commmand line outputs. dispSeg = mynn.DataParallel(dispSeg, cpu_keywords=['im_info', 'roidb'], minibatch=True) ### Training Setups ### args.run_name = misc_utils.get_run_name() output_dir = misc_utils.get_output_dir(args, args.run_name) args.cfg_filename = os.path.basename(args.cfg_file) if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: #from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) ### Training Loop ### dispSeg.train() training_stats = TrainingStats( args, args.disp_interval, tblogger if args.use_tfboard and not args.no_save else None) iters_per_epoch = int(train_size / args.batch_size) # drop last args.iters_per_epoch = iters_per_epoch ckpt_interval_per_epoch = iters_per_epoch // args.ckpt_num_per_epoch try: logger.info('Training starts !') args.step = args.start_iter global_step = iters_per_epoch * args.start_epoch + args.step for args.epoch in range(args.start_epoch, args.start_epoch + args.num_epochs): # ---- Start of epoch ---- # adjust learning rate if args.lr_decay_epochs and args.epoch == args.lr_decay_epochs[ 0] and args.start_iter == 0: args.lr_decay_epochs.pop(0) net_utils.decay_learning_rate(optimizerP, lr, cfg.SOLVER.GAMMA) net_utils.decay_learning_rate(optimizerS, lr, cfg.SOLVER.GAMMA) lr *= cfg.SOLVER.GAMMA for args.step, input_data in zip( range(args.start_iter, iters_per_epoch), dataloader): if cfg.DISP.DISP_ON: input_data['data'] = list( map(lambda x, y: torch.cat((x, y), dim=0), input_data['data'], input_data['data_R'])) if cfg.SEM.DECODER_TYPE.endswith('3Ddeepsup'): input_data['disp_scans'] = torch.arange( 0, cfg.DISP.MAX_DISPLACEMENT).float().view( 1, cfg.DISP.MAX_DISPLACEMENT, 1, 1).repeat(args.batch_size, 1, 1, 1) input_data['semseg_scans'] = torch.arange( 0, cfg.MODEL.NUM_CLASSES).long().view( 1, cfg.MODEL.NUM_CLASSES, 1, 1).repeat(args.batch_size, 1, 1, 1) del input_data['data_R'] for key in input_data: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list( map( lambda x: Variable(x, requires_grad=False).to( 'cuda'), input_data[key])) training_stats.IterTic() net_outputs = dispSeg(**input_data) training_stats.UpdateIterStats(net_outputs) #loss = net_outputs['losses']['loss_semseg'] #acc = net_outputs['metrics']['accuracy_pixel'] #print (loss.item(), acc) #for key in net_outputs.keys(): # print(key) loss = net_outputs['total_loss'] #print("loss.shape:",loss) optimizerP.zero_grad() optimizerS.zero_grad() loss.backward() optimizerP.step() optimizerS.step() training_stats.IterToc() if args.step % args.disp_interval == 0: #disp_image=net_outputs['disp_image'] #semseg_image=net_outputs['semseg_image'] #tblogger.add_image('disp_image',disp_image,global_step) #tblogger.add_image('semseg_image',semseg_image,global_step) log_training_stats(training_stats, global_step, lr) global_step += 1 # ---- End of epoch ---- # save checkpoint net_utils.save_ckpt(output_dir, args, dispSeg, optimizerS) # reset starting iter number after first epoch args.start_iter = 0 # ---- Training ends ---- #if iters_per_epoch % args.disp_interval != 0: # log last stats at the end # log_training_stats(training_stats, global_step, lr) except (RuntimeError, KeyboardInterrupt): logger.info('Save ckpt on exception ...') net_utils.save_ckpt(output_dir, args, dispSeg, optimizerS) logger.info('Save ckpt done.') stack_trace = traceback.format_exc() print(stack_trace) finally: if args.use_tfboard and not args.no_save: tblogger.close()
def main(): """Main function""" args = parse_args() print('Called with args:') print(args) if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") if args.cuda or cfg.NUM_GPUS > 0: cfg.CUDA = True else: raise ValueError("Need Cuda device to run !") if args.dataset == "coco2017": cfg.TRAIN.DATASETS = ('coco_2017_train',) elif args.dataset == "keypoints_coco2017": cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train',) else: raise ValueError("Unexpected args.dataset: {}".format(args.dataset)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) ### Adaptively adjust some configs ### original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH if args.batch_size is None: args.batch_size = original_batch_size cfg.NUM_GPUS = torch.cuda.device_count() assert (args.batch_size % cfg.NUM_GPUS) == 0, \ 'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS) cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS print('Batch size change from {} (in config file) to {}'.format( original_batch_size, args.batch_size)) print('NUM_GPUs: %d, TRAIN.IMS_PER_BATCH: %d' % (cfg.NUM_GPUS, cfg.TRAIN.IMS_PER_BATCH)) if args.num_workers is not None: cfg.DATA_LOADER.NUM_THREADS = args.num_workers print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS) ### Adjust learning based on batch size change linearly old_base_lr = cfg.SOLVER.BASE_LR cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size print('Adjust BASE_LR linearly according to batch size change: {} --> {}'.format( old_base_lr, cfg.SOLVER.BASE_LR)) ### Overwrite some solver settings from command line arguments if args.optimizer is not None: cfg.SOLVER.TYPE = args.optimizer if args.lr is not None: cfg.SOLVER.BASE_LR = args.lr if args.lr_decay_gamma is not None: cfg.SOLVER.GAMMA = args.lr_decay_gamma args.mGPUs = (cfg.NUM_GPUS > 1) timers = defaultdict(Timer) ### Dataset ### timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() train_size = len(roidb) logger.info('{:d} roidb entries'.format(train_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) sampler = MinibatchSampler(ratio_list, ratio_index) dataset = RoiDataLoader( roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, sampler=sampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) assert_and_infer_cfg() ### Model ### maskRCNN = Generalized_RCNN() if cfg.CUDA: maskRCNN.cuda() ### Optimizer ### bias_params = [] nonbias_params = [] for key, value in dict(maskRCNN.named_parameters()).items(): if value.requires_grad: if 'bias' in key: bias_params.append(value) else: nonbias_params.append(value) params = [ {'params': nonbias_params, 'lr': cfg.SOLVER.BASE_LR, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY}, {'params': bias_params, 'lr': cfg.SOLVER.BASE_LR * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0} ] if cfg.SOLVER.TYPE == "SGD": optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizer = torch.optim.Adam(params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) if args.resume: assert checkpoint['iters_per_epoch'] == train_size // args.batch_size, \ "iters_per_epoch should match for resume" # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. # optimizer.load_state_dict(checkpoint['optimizer']) misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) if checkpoint['step'] == (checkpoint['iters_per_epoch'] - 1): # Resume from end of an epoch args.start_epoch = checkpoint['epoch'] + 1 args.start_iter = 0 else: # Resume from the middle of an epoch. # NOTE: dataloader is not synced with previous state args.start_epoch = checkpoint['epoch'] args.start_iter = checkpoint['step'] + 1 del checkpoint torch.cuda.empty_cache() if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values) logging.info("loading Detectron weights %s", args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) lr = optimizer.param_groups[0]['lr'] # lr of non-bias parameters, for commmand line outputs. maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True) ### Training Setups ### run_name = misc_utils.get_run_name() output_dir = misc_utils.get_output_dir(args, run_name) if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) ### Training Loop ### maskRCNN.train() iters_per_epoch = int(train_size / args.batch_size) # drop last ckpt_interval_per_epoch = iters_per_epoch // args.ckpt_num_per_epoch step = 0 try: logger.info('Training starts !') for epoch in range(args.start_epoch, args.start_epoch + args.num_epochs): # ---- Start of epoch ---- loss_avg = 0 timers['train_loop'].tic() # adjust learning rate if args.lr_decay_epochs and epoch == args.lr_decay_epochs[0] and args.start_iter == 0: args.lr_decay_epochs.pop(0) net_utils.decay_learning_rate(optimizer, lr, cfg.SOLVER.GAMMA) lr *= cfg.SOLVER.GAMMA for step, input_data in zip(range(args.start_iter, iters_per_epoch), dataloader): for key in input_data: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list(map(Variable, input_data[key])) outputs = maskRCNN(**input_data) rois_label = outputs['rois_label'] cls_score = outputs['cls_score'] bbox_pred = outputs['bbox_pred'] loss_rpn_cls = outputs['loss_rpn_cls'].mean() loss_rpn_bbox = outputs['loss_rpn_bbox'].mean() loss_rcnn_cls = outputs['loss_rcnn_cls'].mean() loss_rcnn_bbox = outputs['loss_rcnn_bbox'].mean() loss = loss_rpn_cls + loss_rpn_bbox + loss_rcnn_cls + loss_rcnn_bbox if cfg.MODEL.MASK_ON: loss_rcnn_mask = outputs['loss_rcnn_mask'].mean() loss += loss_rcnn_mask if cfg.MODEL.KEYPOINTS_ON: loss_rcnn_keypoints = outputs['loss_rcnn_keypoints'].mean() loss += loss_rcnn_keypoints loss_avg += loss.data.cpu().numpy()[0] optimizer.zero_grad() loss.backward() optimizer.step() if (step+1) % ckpt_interval_per_epoch == 0: net_utils.save_ckpt(output_dir, args, epoch, step, maskRCNN, optimizer, iters_per_epoch) if (step+1) % args.disp_interval == 0: if (step + 1 - args.start_iter) >= args.disp_interval: # for the case of resume diff = timers['train_loop'].toc(average=False) loss_avg /= args.disp_interval loss_rpn_cls = loss_rpn_cls.data[0] loss_rpn_bbox = loss_rpn_bbox.data[0] loss_rcnn_cls = loss_rcnn_cls.data[0] loss_rcnn_bbox = loss_rcnn_bbox.data[0] fg_cnt = torch.sum(rois_label.data.ne(0)) bg_cnt = rois_label.data.numel() - fg_cnt print("[%s][epoch %2d][iter %4d / %4d]" % (run_name, epoch, step, iters_per_epoch)) print("\t\tloss: %.4f, lr: %.2e" % (loss_avg, lr)) print("\t\tfg/bg=(%d/%d), time cost: %f" % (fg_cnt, bg_cnt, diff)) print("\t\trpn_cls: %.4f, rpn_bbox: %.4f, rcnn_cls: %.4f, rcnn_bbox %.4f" % (loss_rpn_cls, loss_rpn_bbox, loss_rcnn_cls, loss_rcnn_bbox)) print_prefix = "\t\t" if cfg.MODEL.MASK_ON: loss_rcnn_mask = loss_rcnn_mask.data[0] print("%srcnn_mask %.4f" % (print_prefix, loss_rcnn_mask)) print_prefix = ", " if cfg.MODEL.KEYPOINTS_ON: loss_rcnn_keypoints = loss_rcnn_keypoints.data[0] print("%srcnn_keypoints %.4f" % (print_prefix, loss_rcnn_keypoints)) if args.use_tfboard: info = { 'loss': loss_avg, 'loss_rpn_cls': loss_rpn_cls, 'loss_rpn_box': loss_rpn_bbox, 'loss_rcnn_cls': loss_rcnn_cls, 'loss_rcnn_box': loss_rcnn_bbox, } if cfg.MODEL.MASK_ON: info['loss_rcnn_mask'] = loss_rcnn_mask if cfg.MODEL.KEYPOINTS_ON: info['loss_rcnn_keypoints'] = loss_rcnn_keypoints for tag, value in info.items(): tblogger.add_scalar(tag, value, iters_per_epoch * epoch + step) loss_avg = 0 timers['train_loop'].tic() # ---- End of epoch ---- # save checkpoint net_utils.save_ckpt(output_dir, args, epoch, step, maskRCNN, optimizer, iters_per_epoch) # reset timer timers['train_loop'].reset() # reset starting iter number after first epoch args.start_iter = 0 except (RuntimeError, KeyboardInterrupt) as e: print('Save on exception:', e) net_utils.save_ckpt(output_dir, args, epoch, step, maskRCNN, optimizer, iters_per_epoch) stack_trace = traceback.format_exc() print(stack_trace) finally: # ---- Training ends ---- if args.use_tfboard: tblogger.close()
def main(): """Main function""" args = parse_cfg() timers = defaultdict(Timer) ### Dataset ### timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) # Effective training sample size for one epoch train_size = roidb_size // args.batch_size * args.batch_size batchSampler = BatchSampler(sampler=MinibatchSampler( ratio_list, ratio_index), batch_size=args.batch_size, drop_last=True) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_sampler=batchSampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) dataiterator = iter(dataloader) ### Model ### maskRCNN = GetRCNNModel() if cfg.CUDA: maskRCNN.cuda() ### Optimizer ### gn_param_nameset = set() for name, module in maskRCNN.named_modules(): if isinstance(module, nn.GroupNorm): gn_param_nameset.add(name + '.weight') gn_param_nameset.add(name + '.bias') gn_params = [] gn_param_names = [] bias_params = [] bias_param_names = [] nonbias_params = [] nonbias_param_names = [] nograd_param_names = [] for key, value in dict(maskRCNN.named_parameters()).items(): if value.requires_grad: if 'bias' in key: bias_params.append(value) bias_param_names.append(key) elif key in gn_param_nameset: gn_params.append(value) gn_param_names.append(key) else: nonbias_params.append(value) nonbias_param_names.append(key) else: nograd_param_names.append(key) assert (gn_param_nameset - set(nograd_param_names) - set(bias_param_names)) == set(gn_param_names) # Learning rate of 0 is a dummy value to be set properly at the start of training params = [{ 'params': nonbias_params, 'lr': 0, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': bias_params, 'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }, { 'params': gn_params, 'lr': 0, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY_GN }] # names of paramerters for each paramter param_names = [nonbias_param_names, bias_param_names, gn_param_names] if cfg.SOLVER.TYPE == "SGD": optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizer = torch.optim.Adam(params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) if args.resume: args.start_step = checkpoint['step'] + 1 if 'train_size' in checkpoint: # For backward compatibility if checkpoint['train_size'] != train_size: print( 'train_size value: %d different from the one in checkpoint: %d' % (train_size, checkpoint['train_size'])) # reorder the params in optimizer checkpoint's params_groups if needed # misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint) # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. # optimizer.load_state_dict(checkpoint['optimizer']) misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) del checkpoint torch.cuda.empty_cache() if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values) logging.info("loading Detectron weights %s", args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) lr = optimizer.param_groups[0][ 'lr'] # lr of non-bias parameters, for commmand line outputs. maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True) if cfg.TRAIN_SYNC_BN: # Shu:For synchorinized BN patch_replication_callback(maskRCNN) ### Training Setups ### args.run_name = misc_utils.get_run_name() + '_step' output_dir = misc_utils.get_output_dir(args, args.run_name) args.cfg_filename = os.path.basename(args.cfg_file) if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) ### Training Loop ### maskRCNN.train() CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS) # Set index for decay steps decay_steps_ind = None for i in range(1, len(cfg.SOLVER.STEPS)): if cfg.SOLVER.STEPS[i] >= args.start_step: decay_steps_ind = i break if decay_steps_ind is None: decay_steps_ind = len(cfg.SOLVER.STEPS) training_stats = TrainingStats( args, args.disp_interval, tblogger if args.use_tfboard and not args.no_save else None) try: logger.info('Training starts !') step = args.start_step for step in range(args.start_step, cfg.SOLVER.MAX_ITER): # Warm up if step < cfg.SOLVER.WARM_UP_ITERS: method = cfg.SOLVER.WARM_UP_METHOD if method == 'constant': warmup_factor = cfg.SOLVER.WARM_UP_FACTOR elif method == 'linear': alpha = step / cfg.SOLVER.WARM_UP_ITERS warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 - alpha) + alpha else: raise KeyError( 'Unknown SOLVER.WARM_UP_METHOD: {}'.format(method)) lr_new = cfg.SOLVER.BASE_LR * warmup_factor net_utils.update_learning_rate(optimizer, lr, lr_new) lr = optimizer.param_groups[0]['lr'] assert lr == lr_new elif step == cfg.SOLVER.WARM_UP_ITERS: net_utils.update_learning_rate(optimizer, lr, cfg.SOLVER.BASE_LR) lr = optimizer.param_groups[0]['lr'] assert lr == cfg.SOLVER.BASE_LR # Learning rate decay if decay_steps_ind < len(cfg.SOLVER.STEPS) and \ step == cfg.SOLVER.STEPS[decay_steps_ind]: logger.info('Decay the learning on step %d', step) lr_new = lr * cfg.SOLVER.GAMMA net_utils.update_learning_rate(optimizer, lr, lr_new) lr = optimizer.param_groups[0]['lr'] assert lr == lr_new decay_steps_ind += 1 training_stats.IterTic() optimizer.zero_grad() for inner_iter in range(args.iter_size): try: input_data = next(dataiterator) except StopIteration: dataiterator = iter(dataloader) input_data = next(dataiterator) for key in input_data: if cfg.MODEL.LR_VIEW_ON or cfg.MODEL.GIF_ON or cfg.MODEL.LRASY_MAHA_ON: if key != 'roidb' and key != 'data': # roidb is a list of ndarrays with inconsistent length input_data[key] = list( map(Variable, input_data[key])) if key == 'data': input_data[key] = [ torch.squeeze(item) for item in input_data[key] ] input_data[key] = list( map(Variable, input_data[key])) else: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list( map(Variable, input_data[key])) net_outputs = maskRCNN(**input_data) training_stats.UpdateIterStats(net_outputs, inner_iter) loss = net_outputs['total_loss'] loss.backward() optimizer.step() training_stats.IterToc() training_stats.LogIterStats(step, lr) if (step + 1) % CHECKPOINT_PERIOD == 0: net_utils.train_save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) # ---- Training ends ---- # Save last checkpoint net_utils.train_save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) except (RuntimeError, KeyboardInterrupt): del dataiterator logger.info('Save ckpt on exception ...') net_utils.train_save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) logger.info('Save ckpt done.') stack_trace = traceback.format_exc() print(stack_trace) finally: if args.use_tfboard and not args.no_save: tblogger.close()
def main(): if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the training code, sry bro :(.") else: cfg.CUDA = True cfg.NUM_GPUS = torch.cuda.device_count() #######~~~.Parameters stuff.~~~####### args = parse_args() print('Called with args:\n', args) # Enables fixed seed if args.fixed_seed: np.random.seed(cfg.RNG_SEED) torch.manual_seed(cfg.RNG_SEED) if cfg.CUDA: torch.cuda.manual_seed_all(cfg.RNG_SEED) torch.backends.cudnn.deterministic = True cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) check_and_overwrite_params(cfg, args) assert_and_infer_cfg() #indentificantion of a specific running args.run_name = misc_utils.get_run_name() + '_step' output_dir = misc_utils.get_output_dir(args) save_training_config(output_dir, cfg, args) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) else: tblogger = None #######~~~.Dataset.~~~####### timers = defaultdict(Timer) # Roi datasets timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) batchSampler = BatchSampler(sampler=MinibatchSampler( ratio_list, ratio_index), batch_size=args.batch_size, drop_last=True) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_sampler=batchSampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) dataiterator = iter(dataloader) ########~~~.Model.~~~####### model = eval(args.model).loot_model(args) if cfg.CUDA: model.cuda() optimizer = OptimizerHandler(model, cfg) load_ckpt(args.load_ckpt, model, optimizer) #######~~~.Training Loop.~~~####### try: # Effective training sample size for one epoch train_size = roidb_size // args.batch_size * args.batch_size CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / (cfg.NUM_GPUS * cfg.TRAIN.ITERATION_SIZE)) training_stats = TrainingStats(args, tblogger) model.train() logger.info('Training starts !') step = args.start_step for step in range(args.start_step, cfg.SOLVER.MAX_ITER): optimizer.update_learning_rate(step) training_stats.IterTic() optimizer.zero_grad() for inner_iter in range(cfg.TRAIN.ITERATION_SIZE): try: input_data = next(dataiterator) except StopIteration: dataiterator = iter(dataloader) input_data = next(dataiterator) for key in input_data: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list(map(Variable, input_data[key])) # input_data['inner_iter'] = torch.tensor((inner_iter)) model.set_inner_iter(step) im_data = input_data['data'][0].cuda() rois = input_data['rois'][0].cuda().type(im_data.dtype) labels = input_data['labels'][0].cuda().type(im_data.dtype) net_outputs = model(im_data, rois, labels) training_stats.UpdateIterStats(net_outputs, inner_iter) loss = net_outputs['total_loss'] loss.backward(retain_graph=True) optimizer.step() training_stats.IterToc() training_stats.LogIterStats(step, optimizer.get_lr()) if (step + 1) % CHECKPOINT_PERIOD == 0: save_ckpt(output_dir, args, step, train_size, model, optimizer) # Training ends, saves the last checkpoint save_ckpt(output_dir, args, step, train_size, model, optimizer) except (RuntimeError, KeyboardInterrupt): del dataiterator logger.info('Save ckpt on exception ...') save_ckpt(output_dir, args, step, train_size, model, optimizer) logger.info('Save ckpt done.') stack_trace = traceback.format_exc() print() print(stack_trace) finally: if args.use_tfboard and not args.no_save: tblogger.close()
def main(): """Main function""" args = parse_args() print('Called with args:') print(args) if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") if args.cuda or cfg.NUM_GPUS > 0: cfg.CUDA = True else: raise ValueError("Need Cuda device to run !") if args.dataset == "coco2017": cfg.TRAIN.DATASETS = ('coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == "keypoints_coco2017": cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 2 # ADE20k as a detection dataset elif args.dataset == "ade_train": cfg.TRAIN.DATASETS = ('ade_train', ) cfg.MODEL.NUM_CLASSES = 446 # Noisy CS6+WIDER datasets elif args.dataset == 'cs6_noise020+WIDER': cfg.TRAIN.DATASETS = ('cs6_noise020', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise050+WIDER': cfg.TRAIN.DATASETS = ('cs6_noise050', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise080+WIDER': cfg.TRAIN.DATASETS = ('cs6_noise080', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise085+WIDER': cfg.TRAIN.DATASETS = ('cs6_noise085', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise090+WIDER': cfg.TRAIN.DATASETS = ('cs6_noise090', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise095+WIDER': cfg.TRAIN.DATASETS = ('cs6_noise095', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise100+WIDER': cfg.TRAIN.DATASETS = ('cs6_noise100', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 # Just Noisy CS6 datasets elif args.dataset == 'cs6_noise020': cfg.TRAIN.DATASETS = ('cs6_noise020', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise030': cfg.TRAIN.DATASETS = ('cs6_noise030', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise040': cfg.TRAIN.DATASETS = ('cs6_noise040', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise050': cfg.TRAIN.DATASETS = ('cs6_noise050', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise060': cfg.TRAIN.DATASETS = ('cs6_noise060', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise070': cfg.TRAIN.DATASETS = ('cs6_noise070', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise080': cfg.TRAIN.DATASETS = ('cs6_noise080', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise085': cfg.TRAIN.DATASETS = ('cs6_noise085', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise090': cfg.TRAIN.DATASETS = ('cs6_noise090', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise095': cfg.TRAIN.DATASETS = ('cs6_noise095', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'cs6_noise100': cfg.TRAIN.DATASETS = ('cs6_noise100', ) cfg.MODEL.NUM_CLASSES = 2 # Cityscapes 7 classes elif args.dataset == "cityscapes": cfg.TRAIN.DATASETS = ('cityscapes_train', ) cfg.MODEL.NUM_CLASSES = 8 # BDD 7 classes elif args.dataset == "bdd_any_any_any": cfg.TRAIN.DATASETS = ('bdd_any_any_any_train', ) cfg.MODEL.NUM_CLASSES = 8 elif args.dataset == "bdd_any_any_daytime": cfg.TRAIN.DATASETS = ('bdd_any_any_daytime_train', ) cfg.MODEL.NUM_CLASSES = 8 elif args.dataset == "bdd_clear_any_daytime": cfg.TRAIN.DATASETS = ('bdd_clear_any_daytime_train', ) cfg.MODEL.NUM_CLASSES = 8 # Cistyscapes Pedestrian sets elif args.dataset == "cityscapes_peds": cfg.TRAIN.DATASETS = ('cityscapes_peds_train', ) cfg.MODEL.NUM_CLASSES = 2 # Cityscapes Car sets elif args.dataset == "cityscapes_cars_HPlen3+kitti_car_train": cfg.TRAIN.DATASETS = ('cityscapes_cars_HPlen3', 'kitti_car_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cityscapes_cars_HPlen5+kitti_car_train": cfg.TRAIN.DATASETS = ('cityscapes_cars_HPlen5', 'kitti_car_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cityscapes_car_train+kitti_car_train": cfg.TRAIN.DATASETS = ('cityscapes_car_train', 'kitti_car_train') cfg.MODEL.NUM_CLASSES = 2 # KITTI Car set elif args.dataset == "kitti_car_train": cfg.TRAIN.DATASETS = ('kitti_car_train', ) cfg.MODEL.NUM_CLASSES = 2 # BDD pedestrians sets elif args.dataset == "bdd_peds": cfg.TRAIN.DATASETS = ('bdd_peds_train', ) # bdd peds: clear_any_daytime cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "bdd_peds_full": cfg.TRAIN.DATASETS = ('bdd_peds_full_train', ) # bdd peds: any_any_any cfg.MODEL.NUM_CLASSES = 2 # Pedestrians with constraints elif args.dataset == "bdd_peds_not_clear_any_daytime": cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train', ) cfg.MODEL.NUM_CLASSES = 2 # Ashish's 20k samples videos elif args.dataset == "bdd_peds_not_clear_any_daytime_20k": cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_20k_train', ) cfg.MODEL.NUM_CLASSES = 2 # Source domain + Target domain detections elif args.dataset == "bdd_peds+DETS_20k": cfg.TRAIN.DATASETS = ('bdd_peds_dets_20k_target_domain', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 # Source domain + Target domain detections -- same 18k images as HP18k elif args.dataset == "bdd_peds+DETS18k": cfg.TRAIN.DATASETS = ('bdd_peds_dets18k_target_domain', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 # Only Dets elif args.dataset == "DETS20k": cfg.TRAIN.DATASETS = ('bdd_peds_dets_20k_target_domain', ) cfg.MODEL.NUM_CLASSES = 2 # Only Dets18k - same images as HP18k elif args.dataset == 'DETS18k': cfg.TRAIN.DATASETS = ('bdd_peds_dets18k_target_domain', ) cfg.MODEL.NUM_CLASSES = 2 # Only HP elif args.dataset == 'HP': cfg.TRAIN.DATASETS = ('bdd_peds_HP_target_domain', ) cfg.MODEL.NUM_CLASSES = 2 # Only HP 18k videos elif args.dataset == 'HP18k': cfg.TRAIN.DATASETS = ('bdd_peds_HP18k_target_domain', ) cfg.MODEL.NUM_CLASSES = 2 # Source domain + Target domain HP elif args.dataset == 'bdd_peds+HP': cfg.TRAIN.DATASETS = ('bdd_peds_train', 'bdd_peds_HP_target_domain') cfg.MODEL.NUM_CLASSES = 2 # Source domain + Target domain HP 18k videos cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k': cfg.TRAIN.DATASETS = ('bdd_peds_HP18k_target_domain', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 #### Source domain + Target domain with different conf threshold theta #### elif args.dataset == 'bdd_peds+HP18k_thresh-050': cfg.TRAIN.DATASETS = ('bdd_HP18k_thresh-050', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k_thresh-060': cfg.TRAIN.DATASETS = ('bdd_HP18k_thresh-060', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k_thresh-070': cfg.TRAIN.DATASETS = ('bdd_HP18k_thresh-070', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k_thresh-090': cfg.TRAIN.DATASETS = ('bdd_HP18k_thresh-090', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 ############################## #### Data distillation on BDD -- for rebuttal elif args.dataset == 'bdd_peds+bdd_data_dist_small': cfg.TRAIN.DATASETS = ('bdd_data_dist_small', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+bdd_data_dist_mid': cfg.TRAIN.DATASETS = ('bdd_data_dist_mid', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+bdd_data_dist': cfg.TRAIN.DATASETS = ('bdd_data_dist', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 ############################## #### Source domain + **Labeled** Target domain with varying number of images elif args.dataset == 'bdd_peds+labeled_100': cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_100', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+labeled_075': cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_075', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+labeled_050': cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_050', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+labeled_025': cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_025', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+labeled_010': cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_010', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+labeled_005': cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_005', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+labeled_001': cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_001', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 ############################## # Source domain + Target domain HP tracker bboxes only elif args.dataset == 'bdd_peds+HP18k_track_only': cfg.TRAIN.DATASETS = ('bdd_HP18k_track_only', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 ##### subsets of bdd_HP18k with different constraints # Source domain + HP tracker images at NIGHT elif args.dataset == 'bdd_peds+HP18k_any_any_night': cfg.TRAIN.DATASETS = ('bdd_HP18k_any_any_night', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k_rainy_any_daytime': cfg.TRAIN.DATASETS = ('bdd_HP18k_rainy_any_daytime', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k_rainy_any_night': cfg.TRAIN.DATASETS = ('bdd_HP18k_rainy_any_night', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k_overcast,rainy_any_daytime': cfg.TRAIN.DATASETS = ('bdd_HP18k_overcast,rainy_any_daytime', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k_overcast,rainy_any_night': cfg.TRAIN.DATASETS = ('bdd_HP18k_overcast,rainy_any_night', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k_overcast,rainy,snowy_any_daytime': cfg.TRAIN.DATASETS = ('bdd_HP18k_overcast,rainy,snowy_any_daytime', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 ############# end of bdd constraned subsets ##################### # Source domain + Target domain HP18k -- after histogram matching elif args.dataset == 'bdd_peds+HP18k_remap_hist': cfg.TRAIN.DATASETS = ('bdd_peds_HP18k_target_domain_remap_hist', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+HP18k_remap_cityscape_hist': cfg.TRAIN.DATASETS = ( 'bdd_peds_HP18k_target_domain_remap_cityscape_hist', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 # Source domain + Target domain HP18k -- after histogram matching elif args.dataset == 'bdd_peds+HP18k_remap_random': cfg.TRAIN.DATASETS = ('bdd_peds_HP18k_target_domain_remap_random', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 # Source+Noisy Target domain -- prevent domain adv from using HP roi info elif args.dataset == 'bdd_peds+bdd_HP18k_noisy_100k': cfg.TRAIN.DATASETS = ('bdd_HP18k_noisy_100k', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+bdd_HP18k_noisy_080': cfg.TRAIN.DATASETS = ('bdd_HP18k_noisy_080', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+bdd_HP18k_noisy_060': cfg.TRAIN.DATASETS = ('bdd_HP18k_noisy_060', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == 'bdd_peds+bdd_HP18k_noisy_070': cfg.TRAIN.DATASETS = ('bdd_HP18k_noisy_070', 'bdd_peds_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "wider_train": cfg.TRAIN.DATASETS = ('wider_train', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-subset": cfg.TRAIN.DATASETS = ('cs6-subset', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-subset-score": cfg.TRAIN.DATASETS = ('cs6-subset-score', ) elif args.dataset == "cs6-subset-gt": cfg.TRAIN.DATASETS = ('cs6-subset-gt', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-3013-gt": cfg.TRAIN.DATASETS = ('cs6-3013-gt', ) # DEBUG: overfit on one video annots cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-subset-gt+WIDER": cfg.TRAIN.DATASETS = ('cs6-subset-gt', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-subset+WIDER": cfg.TRAIN.DATASETS = ('cs6-subset', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-gt": cfg.TRAIN.DATASETS = ('cs6-train-gt', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-gt-noisy-0.3": cfg.TRAIN.DATASETS = ('cs6-train-gt-noisy-0.3', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-gt-noisy-0.5": cfg.TRAIN.DATASETS = ('cs6-train-gt-noisy-0.5', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-det-score": cfg.TRAIN.DATASETS = ('cs6-train-det-score', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-det-score-0.5": cfg.TRAIN.DATASETS = ('cs6-train-det-score-0.5', ) elif args.dataset == "cs6-train-det": cfg.TRAIN.DATASETS = ('cs6-train-det', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-det-0.5": cfg.TRAIN.DATASETS = ('cs6-train-det-0.5', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-hp": cfg.TRAIN.DATASETS = ('cs6-train-hp', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-easy-gt": cfg.TRAIN.DATASETS = ('cs6-train-easy-gt', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-easy-gt-sub": cfg.TRAIN.DATASETS = ('cs6-train-easy-gt-sub', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-easy-hp": cfg.TRAIN.DATASETS = ('cs6-train-easy-hp', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-easy-det": cfg.TRAIN.DATASETS = ('cs6-train-easy-det', ) cfg.MODEL.NUM_CLASSES = 2 # Joint training with CS6 and WIDER elif args.dataset == "cs6-train-easy-gt-sub+WIDER": cfg.TRAIN.DATASETS = ('cs6-train-easy-gt-sub', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-gt+WIDER": cfg.TRAIN.DATASETS = ('cs6-train-gt', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-hp+WIDER": cfg.TRAIN.DATASETS = ('cs6-train-hp', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-dummy+WIDER": cfg.TRAIN.DATASETS = ('cs6-train-dummy', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "cs6-train-det+WIDER": cfg.TRAIN.DATASETS = ('cs6-train-det', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 # Dets dataset created by removing tracker results from the HP json elif args.dataset == "cs6_train_det_from_hp+WIDER": cfg.TRAIN.DATASETS = ('cs6_train_det_from_hp', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 # Dataset created by removing det results from the HP json -- HP tracker only elif args.dataset == "cs6_train_hp_tracker_only+WIDER": cfg.TRAIN.DATASETS = ('cs6_train_hp_tracker_only', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 # HP dataset with noisy labels: used to prevent DA from getting any info from HP elif args.dataset == "cs6_train_hp_noisy_100+WIDER": cfg.TRAIN.DATASETS = ('cs6_train_hp_noisy_100', 'wider_train') cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError("Unexpected args.dataset: {}".format(args.dataset)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) ### Adaptively adjust some configs ### original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH original_num_gpus = cfg.NUM_GPUS if args.batch_size is None: args.batch_size = original_batch_size cfg.NUM_GPUS = torch.cuda.device_count() assert (args.batch_size % cfg.NUM_GPUS) == 0, \ 'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS) cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS effective_batch_size = args.iter_size * args.batch_size print('effective_batch_size = batch_size * iter_size = %d * %d' % (args.batch_size, args.iter_size)) print('Adaptive config changes:') print(' effective_batch_size: %d --> %d' % (original_batch_size, effective_batch_size)) print(' NUM_GPUS: %d --> %d' % (original_num_gpus, cfg.NUM_GPUS)) print(' IMS_PER_BATCH: %d --> %d' % (original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH)) ### Adjust learning based on batch size change linearly # For iter_size > 1, gradients are `accumulated`, so lr is scaled based # on batch_size instead of effective_batch_size old_base_lr = cfg.SOLVER.BASE_LR cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size print('Adjust BASE_LR linearly according to batch_size change:\n' ' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR)) ### Adjust solver steps step_scale = original_batch_size / effective_batch_size old_solver_steps = cfg.SOLVER.STEPS old_max_iter = cfg.SOLVER.MAX_ITER cfg.SOLVER.STEPS = list( map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS)) cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5) print( 'Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n' ' SOLVER.STEPS: {} --> {}\n' ' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps, cfg.SOLVER.STEPS, old_max_iter, cfg.SOLVER.MAX_ITER)) # Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function # of `collect_and_distribute_fpn_rpn_proposals.py` # # post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5) if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN: cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch print( 'Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n' ' cfg.FPN.RPN_COLLECT_SCALE: {}'.format( cfg.FPN.RPN_COLLECT_SCALE)) if args.num_workers is not None: cfg.DATA_LOADER.NUM_THREADS = args.num_workers print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS) ### Overwrite some solver settings from command line arguments if args.optimizer is not None: cfg.SOLVER.TYPE = args.optimizer if args.lr is not None: cfg.SOLVER.BASE_LR = args.lr if args.lr_decay_gamma is not None: cfg.SOLVER.GAMMA = args.lr_decay_gamma assert_and_infer_cfg() timers = defaultdict(Timer) """def _init_fn(worker_id): random.seed(999) np.random.seed(999) torch.cuda.manual_seed(999) torch.cuda.manual_seed_all(999) torch.manual_seed(999) torch.backends.cudnn.deterministic = True """ ### Dataset ### timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) if cfg.TRAIN.JOINT_TRAINING: if len(cfg.TRAIN.DATASETS) == 2: print('Joint training on two datasets') else: raise NotImplementedError joint_training_roidb = [] for i, dataset_name in enumerate(cfg.TRAIN.DATASETS): # ROIDB construction timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( (dataset_name), cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) if i == 0: roidb_size = len(roidb) batchSampler = BatchSampler(sampler=MinibatchSampler( ratio_list, ratio_index), batch_size=args.batch_size, drop_last=True) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_sampler=batchSampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) #worker_init_fn=_init_fn) # decrease num-threads when using two dataloaders dataiterator = iter(dataloader) joint_training_roidb.append({ 'dataloader': dataloader, 'dataiterator': dataiterator, 'dataset_name': dataset_name }) else: roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) batchSampler = BatchSampler(sampler=MinibatchSampler( ratio_list, ratio_index), batch_size=args.batch_size, drop_last=True) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_sampler=batchSampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) #worker_init_fn=init_fn) dataiterator = iter(dataloader) # Effective training sample size for one epoch train_size = roidb_size // args.batch_size * args.batch_size if cfg.TRAIN.JOINT_SELECTIVE_FG: orig_fg_batch_ratio = cfg.TRAIN.FG_FRACTION ### Model ### maskRCNN = Generalized_RCNN() if cfg.CUDA: maskRCNN.cuda() ### Optimizer ### gn_param_nameset = set() for name, module in maskRCNN.named_modules(): if isinstance(module, nn.GroupNorm): gn_param_nameset.add(name + '.weight') gn_param_nameset.add(name + '.bias') gn_params = [] gn_param_names = [] bias_params = [] bias_param_names = [] nonbias_params = [] nonbias_param_names = [] nograd_param_names = [] for key, value in dict(maskRCNN.named_parameters()).items(): if value.requires_grad: if 'bias' in key: bias_params.append(value) bias_param_names.append(key) elif key in gn_param_nameset: gn_params.append(value) gn_param_names.append(key) else: nonbias_params.append(value) nonbias_param_names.append(key) else: nograd_param_names.append(key) assert (gn_param_nameset - set(nograd_param_names) - set(bias_param_names)) == set(gn_param_names) # Learning rate of 0 is a dummy value to be set properly at the start of training params = [{ 'params': nonbias_params, 'lr': 0, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': bias_params, 'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }, { 'params': gn_params, 'lr': 0, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY_GN }] # names of paramerters for each paramter param_names = [nonbias_param_names, bias_param_names, gn_param_names] if cfg.SOLVER.TYPE == "SGD": optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizer = torch.optim.Adam(params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) if args.resume: args.start_step = checkpoint['step'] + 1 if 'train_size' in checkpoint: # For backward compatibility if checkpoint['train_size'] != train_size: print( 'train_size value: %d different from the one in checkpoint: %d' % (train_size, checkpoint['train_size'])) # reorder the params in optimizer checkpoint's params_groups if needed # misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint) # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. # optimizer.load_state_dict(checkpoint['optimizer']) misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) del checkpoint torch.cuda.empty_cache() if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values) logging.info("loading Detectron weights %s", args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) lr = optimizer.param_groups[0][ 'lr'] # lr of non-bias parameters, for commmand line outputs. maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True) ### Training Setups ### args.run_name = misc_utils.get_run_name() + '_' + str( cfg.TRAIN.DATASETS) + '_step' output_dir = misc_utils.get_output_dir(args, args.run_name) args.cfg_filename = os.path.basename(args.cfg_file) if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) ### Training Loop ### maskRCNN.train() CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS) # Set index for decay steps decay_steps_ind = None for i in range(1, len(cfg.SOLVER.STEPS)): if cfg.SOLVER.STEPS[i] >= args.start_step: decay_steps_ind = i break if decay_steps_ind is None: decay_steps_ind = len(cfg.SOLVER.STEPS) training_stats = TrainingStats( args, args.disp_interval, tblogger if args.use_tfboard and not args.no_save else None) try: logger.info('Training starts !') step = args.start_step for step in range(args.start_step, cfg.SOLVER.MAX_ITER): """ random.seed(cfg.RNG_SEED) np.random.seed(cfg.RNG_SEED) torch.cuda.manual_seed(cfg.RNG_SEED) torch.cuda.manual_seed_all(cfg.RNG_SEED) torch.manual_seed(cfg.RNG_SEED) torch.backends.cudnn.deterministic = True """ # Warm up if step < cfg.SOLVER.WARM_UP_ITERS: method = cfg.SOLVER.WARM_UP_METHOD if method == 'constant': warmup_factor = cfg.SOLVER.WARM_UP_FACTOR elif method == 'linear': alpha = step / cfg.SOLVER.WARM_UP_ITERS warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 - alpha) + alpha else: raise KeyError( 'Unknown SOLVER.WARM_UP_METHOD: {}'.format(method)) lr_new = cfg.SOLVER.BASE_LR * warmup_factor net_utils.update_learning_rate(optimizer, lr, lr_new) lr = optimizer.param_groups[0]['lr'] assert lr == lr_new elif step == cfg.SOLVER.WARM_UP_ITERS: net_utils.update_learning_rate(optimizer, lr, cfg.SOLVER.BASE_LR) lr = optimizer.param_groups[0]['lr'] assert lr == cfg.SOLVER.BASE_LR # Learning rate decay if decay_steps_ind < len(cfg.SOLVER.STEPS) and \ step == cfg.SOLVER.STEPS[decay_steps_ind]: logger.info('Decay the learning on step %d', step) lr_new = lr * cfg.SOLVER.GAMMA net_utils.update_learning_rate(optimizer, lr, lr_new) lr = optimizer.param_groups[0]['lr'] assert lr == lr_new decay_steps_ind += 1 training_stats.IterTic() optimizer.zero_grad() for inner_iter in range(args.iter_size): # use a iter counter for optional alternating batches if args.iter_size == 1: iter_counter = step else: iter_counter = inner_iter if cfg.TRAIN.JOINT_TRAINING: # alternate batches between dataset[0] and dataset[1] if iter_counter % 2 == 0: if True: #DEBUG: print('Dataset: %s' % joint_training_roidb[0]['dataset_name']) dataloader = joint_training_roidb[0]['dataloader'] dataiterator = joint_training_roidb[0]['dataiterator'] # NOTE: if available FG samples cannot fill minibatch # then batchsize will be smaller than cfg.TRAIN.BATCH_SIZE_PER_IM. else: if True: #DEBUG: print('Dataset: %s' % joint_training_roidb[1]['dataset_name']) dataloader = joint_training_roidb[1]['dataloader'] dataiterator = joint_training_roidb[1]['dataiterator'] try: input_data = next(dataiterator) except StopIteration: # end of epoch for dataloader dataiterator = iter(dataloader) input_data = next(dataiterator) if cfg.TRAIN.JOINT_TRAINING: if iter_counter % 2 == 0: joint_training_roidb[0][ 'dataiterator'] = dataiterator else: joint_training_roidb[1][ 'dataiterator'] = dataiterator for key in input_data: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list(map(Variable, input_data[key])) net_outputs = maskRCNN(**input_data) training_stats.UpdateIterStats(net_outputs, inner_iter) loss = net_outputs['total_loss'] # [p.data.get_device() for p in maskRCNN.parameters()] # [(name, p.data.get_device()) for name, p in maskRCNN.named_parameters()] loss.backward() optimizer.step() training_stats.IterToc() training_stats.LogIterStats(step, lr) if (step + 1) % CHECKPOINT_PERIOD == 0: save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) # ---- Training ends ---- # Save last checkpoint save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) except (RuntimeError, KeyboardInterrupt): del dataiterator logger.info('Save ckpt on exception ...') save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) logger.info('Save ckpt done.') stack_trace = traceback.format_exc() print(stack_trace) finally: if args.use_tfboard and not args.no_save: tblogger.close()
def main(): """Main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") logger = utils.logging.setup_logging(__name__) args = parse_args() logger.info('Called with args:') logger.info(args) assert (torch.cuda.device_count() == 1) ^ bool(args.multi_gpu_testing) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' if args.output_dir is None: ckpt_path = args.load_ckpt if args.load_ckpt else args.load_detectron args.output_dir = os.path.join( os.path.dirname(os.path.dirname(ckpt_path)), 'test') logger.info('Automatically set output directory to %s', args.output_dir) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) if args.close_fpn: args.cfg_file = "configs/few_shot/e2e_mask_rcnn_R-50-C4_1x_{}.yaml".format(args.group) else: args.cfg_file = "configs/few_shot/e2e_mask_rcnn_R-50-FPN_1x_{}.yaml".format(args.group) if args.cfg_file is not None: merge_cfg_from_file(args.cfg_file) if args.set_cfgs is not None: merge_cfg_from_list(args.set_cfgs) cfg.VIS = args.vis cfg.SEEN = args.seen if args.close_co_atten: cfg.CO_ATTEN = False if args.close_relation_rcnn: cfg.RELATION_RCNN = False if not args.close_fpn: cfg.FAST_RCNN.ROI_BOX_HEAD = 'fast_rcnn_heads.roi_2mlp_head' cfg.MRCNN.ROI_MASK_HEAD = 'mask_rcnn_heads.mask_rcnn_fcn_head_v1up4convs' else: cfg.FAST_RCNN.ROI_BOX_HEAD = 'torchResNet.ResNet_roi_conv5_head' cfg.MRCNN.ROI_MASK_HEAD = 'mask_rcnn_heads.mask_rcnn_fcn_head_v0upshare' if args.deform_conv: cfg.MODEL.USE_DEFORM = True if args.dataset == "fis_cell": cfg.TEST.DATASETS = ('fis_cell_test',) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "coco2017": cfg.TEST.DATASETS = ('coco_2017_val',) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == "keypoints_coco2017": cfg.TEST.DATASETS = ('keypoints_coco_2017_val',) cfg.MODEL.NUM_CLASSES = 2 else: # For subprocess call assert cfg.TEST.DATASETS, 'cfg.TEST.DATASETS shouldn\'t be empty' assert_and_infer_cfg() #logger.info('Testing with config:') #logger.info(pprint.pformat(cfg)) # manually set args.cuda args.cuda = True timer_for_ds = defaultdict(Timer) ### Dataset ### timer_for_ds['roidb'].tic() imdb, roidb, ratio_list, ratio_index, query, cat_list = combined_roidb( cfg.TEST.DATASETS, False) timer_for_ds['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timer_for_ds['roidb'].average_time) batchSampler = BatchSampler( sampler=MinibatchSampler(ratio_list, ratio_index, shuffle=False), batch_size=1, drop_last=False ) dataset = RoiDataLoader( roidb, ratio_list, ratio_index, query, cfg.MODEL.NUM_CLASSES, training=False, cat_list=cat_list, shot=args.checkshot) ### Model ### model = initialize_model_from_cfg(args, gpu_id=0) all_results = OrderedDict({ 'box': OrderedDict( [ ('AP', []), ('AP50', []), ('AP75', []), ('APs', []), ('APm', []), ('APl', []), ] ), 'mask': OrderedDict( [ ('AP', []), ('AP50', []), ('AP75', []), ('APs', []), ('APm', []), ('APl', []), ] ) }) timer_for_total = defaultdict(Timer) timer_for_total['total_test_time'].tic() for avg in range(args.average): dataset.query_position = avg dataloader = torch.utils.data.DataLoader( dataset, batch_sampler=batchSampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch ) dataiterator = iter(dataloader) num_images = len(ratio_index) num_cats = imdb.num_classes all_boxes, all_segms, all_keyps = empty_results(num_cats, num_images) # total quantity of testing images num_detect = len(ratio_index) timers = defaultdict(Timer) post_fix = '%dshot_g%d_seen%d_%d'%(args.checkshot, args.group, args.seen, avg) for i, index in enumerate(ratio_index): input_data = next(dataiterator) catgory = input_data['choice'] entry = dataset._roidb[dataset.ratio_index[i]] if cfg.TEST.PRECOMPUTED_PROPOSALS: # The roidb may contain ground-truth rois (for example, if the roidb # comes from the training or val split). We only want to evaluate # detection on the *non*-ground-truth rois. We select only the rois # that have the gt_classes field set to 0, which means there's no # ground truth. box_proposals = entry['boxes'][entry['gt_classes'] == 0] if len(box_proposals) == 0: continue else: # Faster R-CNN type models generate proposals on-the-fly with an # in-network RPN; 1-stage models don't require proposals. box_proposals = None #im = cv2.imread(entry['image']) im = imread(entry['image']) if len(im.shape) == 2: im = im[:,:,np.newaxis] im = np.concatenate((im,im,im), axis=2) alpha = 2.2 # Simple contrast control beta = 0.5 # Simple brightness control im_colored = im.copy() if catgory[0].item() == 7: im_colored[:,:,0] = 0 im_colored[:,:,1] = 0 im_colored = cv2.convertScaleAbs(im_colored, alpha=alpha, beta=beta) elif catgory[0].item() == 8: im_colored[:,:,1] = 0 im_colored[:,:,2] = 0 im_colored = cv2.convertScaleAbs(im_colored, alpha=alpha, beta=beta) elif catgory[0].item() == 6: im_colored[:,:,0] = 0 im_colored[:,:,2] = 0 im_colored = cv2.convertScaleAbs(im_colored, alpha=alpha, beta=beta) cls_boxes_i, cls_segms_i, cls_keyps_i = im_detect_all(model, im, input_data['query'], input_data['query_type'], catgory, num_cats, box_proposals, timers) im = im_colored extend_results(i, all_boxes, cls_boxes_i) if cls_segms_i is not None: extend_results(i, all_segms, cls_segms_i) if cls_keyps_i is not None: extend_results(i, all_keyps, cls_keyps_i) if i % 10 == 0: # Reduce log file size ave_total_time = np.sum([t.average_time for t in timers.values()]) eta_seconds = ave_total_time * (num_images - i - 1) eta = str(datetime.timedelta(seconds=int(eta_seconds))) det_time = ( timers['im_detect_bbox'].average_time + timers['im_detect_mask'].average_time + timers['im_detect_keypoints'].average_time ) misc_time = ( timers['misc_bbox'].average_time + timers['misc_mask'].average_time + timers['misc_keypoints'].average_time ) logger.info( ( 'im_detect: range [{:d}, {:d}] of {:d}: ' '{:d}/{:d} {:.3f}s + {:.3f}s (eta: {})' ).format( 1, num_detect, num_detect, i + 1, num_detect, det_time, misc_time, eta ) ) if cfg.VIS: im_name = entry['image'] class_name = im_name.split('/')[-4] file_name = im_name.split('/')[-3] im_target = im.copy() to_tensor = transforms.ToTensor() o_querys=[] for i in range(args.checkshot): o_query = input_data['query'][0][i][0].permute(1, 2,0).contiguous().cpu().numpy() o_query *= [0.229, 0.224, 0.225] o_query += [0.485, 0.456, 0.406] o_query *= 255 o_query_colored = o_query.copy() if catgory[0].item() == 7: o_query_colored[:,:,0] = 0 o_query_colored[:,:,1] = 0 o_query_colored = cv2.convertScaleAbs(o_query_colored, alpha=alpha, beta=beta) elif catgory[0].item() == 8: o_query_colored[:,:,1] = 0 o_query_colored[:,:,2] = 0 o_query_colored = cv2.convertScaleAbs(o_query_colored, alpha=alpha, beta=beta) elif catgory[0].item() == 6: o_query_colored[:,:,0] = 0 o_query_colored[:,:,2] = 0 o_query_colored = cv2.convertScaleAbs(o_query_colored, alpha=alpha, beta=beta) o_query = o_query_colored o_query = Image.fromarray(o_query.astype(np.uint8)) o_querys.append(to_tensor(o_query)) o_querys_grid = make_grid(o_querys, nrow=args.checkshot//2, normalize=True, scale_each=True, pad_value=1) o_querys_grid = transforms.ToPILImage()(o_querys_grid).convert("RGB") query_w, query_h = o_querys_grid.size query_bg = Image.new('RGB', (im_target.shape[1], im_target.shape[0]), (255, 255, 255)) bg_w, bg_h = query_bg.size offset = ((bg_w - query_w) // 2, (bg_h - query_h) // 2) query_bg.paste(o_querys_grid, offset) query = np.asarray(query_bg) im_pair = np.concatenate((im_target, query), axis=1) im_output_dir = os.path.join(args.output_dir, 'vis', post_fix, class_name) if not os.path.exists(im_output_dir): os.makedirs(im_output_dir) sample_output_dir = os.path.join(im_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name))) if not os.path.exists(sample_output_dir): os.makedirs(sample_output_dir) target_save_name = os.path.join(sample_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name)) + '_target.pdf') target = Image.fromarray(im_target.astype(np.uint8)) target.save(target_save_name,"pdf") query_save_name = os.path.join(sample_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name)) + '_query.pdf') query = Image.fromarray(query.astype(np.uint8)) query.save(query_save_name,"pdf") pred_save_name = os.path.join(sample_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name)) + '_pred.pdf') vis_utils.save_one_image( im, pred_save_name, cls_boxes_i, segms = cls_segms_i, keypoints = cls_keyps_i, thresh = cfg.VIS_TH, box_alpha = 0.6, dataset = imdb, show_class = False ) im_det = vis_utils.vis_one_image( im, '{:d}_det_{:s}'.format(i, file_name), os.path.join(args.output_dir, 'vis', post_fix), cls_boxes_i, segms = cls_segms_i, keypoints = cls_keyps_i, thresh = cfg.VIS_TH, box_alpha = 0.6, dataset = imdb, show_class = False, class_name = class_name, draw_bbox = False ) gt_save_name = os.path.join(sample_output_dir, os.path.basename('{:d}_{:s}'.format(i, file_name)) + '_gt.pdf') vis_utils.save_one_image_gt( im, entry['id'], gt_save_name, dataset = imdb) im_gt = vis_utils.vis_one_image_gt( im, entry['id'], '{:d}_gt_{:s}'.format(i, file_name), os.path.join(args.output_dir, 'vis', post_fix), dataset = imdb, class_name = class_name) im_det = np.asarray(im_det) im_gt = np.asarray(im_gt) im2draw = np.concatenate((im_gt, im_det), axis=1) im2show = np.concatenate((im_pair, im2draw), axis=0) im_save_name = os.path.basename('{:d}_{:s}'.format(i, file_name)) + '.png' cv2.imwrite(os.path.join(im_output_dir, '{}'.format(im_save_name)), cv2.cvtColor(im2show, cv2.COLOR_RGB2BGR)) """ if cfg.VIS: im_name = dataset._roidb[dataset.ratio_index[i]]['image'] class_name = im_name.split('/')[-4] file_name = im_name.split('/')[-3] im2show = Image.open(im_name).convert("RGB") o_query = input_data['query'][0][0][0].permute(1, 2, 0).contiguous().cpu().numpy() o_query *= [0.229, 0.224, 0.225] o_query += [0.485, 0.456, 0.406] o_query *= 255 o_query = o_query[:,:,::-1] o_query = Image.fromarray(o_query.astype(np.uint8)) query_w, query_h = o_query.size query_bg = Image.new('RGB', (im2show.size), (255, 255, 255)) bg_w, bg_h = query_bg.size offset = ((bg_w - query_w) // 2, (bg_h - query_h) // 2) query_bg.paste(o_query, offset) im2show = np.asarray(im2show) o_query = np.asarray(query_bg) im2show = np.concatenate((im2show, o_query), axis=1) output_dir = os.path.join(args.output_dir, 'vis') if not os.path.exists(output_dir): os.makedirs(output_dir) im_save_dir = os.path.join(output_dir, class_name) if not os.path.exists(im_save_dir): os.makedirs(im_save_dir) im_save_name = os.path.join(im_save_dir, file_name + '_%d_d.png'%(i)) cv2.imwrite(im_save_name, cv2.cvtColor(im2show, cv2.COLOR_RGB2BGR)) """ cfg_yaml = yaml.dump(cfg) #det_name = 'detections.pkl' det_file = os.path.join(args.output_dir, 'detections_' + post_fix + '.pkl') save_object( dict( all_boxes=all_boxes, all_segms=all_segms, all_keyps=all_keyps, cfg=cfg_yaml ), det_file ) logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file))) results = task_evaluation.evaluate_all( imdb, all_boxes, all_segms, all_keyps, args.output_dir ) task_evaluation.check_expected_results( results, atol=cfg.EXPECTED_RESULTS_ATOL, rtol=cfg.EXPECTED_RESULTS_RTOL ) for task, metrics in all_results.items(): metric_names = metrics.keys() for metric_name in metric_names: all_results[task][metric_name].append(results[imdb.name][task][metric_name]) #task_evaluation.log_copy_paste_friendly_results(results) for task, metrics in all_results.items(): metric_names = metrics.keys() for metric_name in metric_names: values = all_results[task][metric_name] all_results[task][metric_name] = sum(values) / len(values) post_fix = '%dshot_g%d_seen%d'%(args.checkshot, args.group, args.seen) avg_results_path = os.path.join(args.output_dir, ('avg_cocoeval_' + post_fix + '_results.json')) with open(avg_results_path, 'w') as f: f.write(json.dumps(all_results)) timer_for_total['total_test_time'].toc() logger.info('Total inference time: {:.3f}s'.format(timer_for_total['total_test_time'].average_time))
def main(): """Main function""" args = parse_args() print('Called with args:') print(args) if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") if args.cuda or cfg.NUM_GPUS > 0: cfg.CUDA = True else: raise ValueError("Need Cuda device to run !") if args.dataset == "coco2017": cfg.TRAIN.DATASETS = ('coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == "coco2014": cfg.TRAIN.DATASETS = ('coco_2014_train', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == 'vcoco_trainval': cfg.TRAIN.DATASETS = ('vcoco_trainval', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == 'vcoco_train': cfg.TRAIN.DATASETS = ('vcoco_train', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == 'vcoco_val': cfg.TRAIN.DATASETS = ('vcoco_val', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == 'keypoints_coco2014': cfg.TRAIN.DATASETS = ('keypoints_coco_2014_train', ) cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "keypoints_coco2017": cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError("Unexpected args.dataset: {}".format(args.dataset)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) if args.vcoco_kp_on: cfg.VCOCO.KEYPOINTS_ON = True cfg.NETWORK_NAME = args.net_name # network name print('Network name:', args.net_name) cfg.MODEL.CONV_BODY = args.conv_body # backbone network name print('Conv_body name:', args.conv_body) cfg.TRAIN.FG_THRESH = args.fg_thresh print('Train fg thresh:', args.fg_thresh) cfg.RESNETS.FREEZE_AT = args.freeze_at print('Freeze at: ', args.freeze_at) cfg.VCOCO.MLP_HEAD_DIM = args.mlp_head_dim print('MLP head dim: ', args.mlp_head_dim) cfg.SOLVER.MAX_ITER = args.max_iter print('MAX iter: ', args.max_iter) cfg.TRAIN.SNAPSHOT_ITERS = args.snapshot print('Snapshot Iters: ', args.snapshot) if args.solver_steps is not None: cfg.SOLVER.STEPS = args.solver_steps print('Solver_steps: ', cfg.SOLVER.STEPS) cfg.VCOCO.TRIPLETS_NUM_PER_IM = args.triplets_num_per_im print('triplets_num_per_im: ', cfg.VCOCO.TRIPLETS_NUM_PER_IM) cfg.VCOCO.HEATMAP_KERNEL_SIZE = args.heatmap_kernel_size print('heatmap_kernel_size: ', cfg.VCOCO.HEATMAP_KERNEL_SIZE) cfg.VCOCO.PART_CROP_SIZE = args.part_crop_size print('part_crop_size: ', cfg.VCOCO.PART_CROP_SIZE) print('use use_kps17 for part Align: ', args.use_kps17) if args.use_kps17: cfg.VCOCO.USE_KPS17 = True else: cfg.VCOCO.USE_KPS17 = False print('MULTILEVEL_ROIS: ', cfg.FPN.MULTILEVEL_ROIS) if args.vcoco_use_spatial: cfg.VCOCO.USE_SPATIAL = True if args.vcoco_use_union_feat: cfg.VCOCO.USE_UNION_FEAT = True if args.use_precomp_box: cfg.VCOCO.USE_PRECOMP_BOX = True cfg.DEBUG_TEST_WITH_GT = True if args.lr is not None: cfg.SOLVER.BASE_LR = args.lr ### Adaptively adjust some configs ### original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH # 16 original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH original_num_gpus = cfg.NUM_GPUS if args.batch_size is None: args.batch_size = original_batch_size cfg.NUM_GPUS = torch.cuda.device_count() assert (args.batch_size % cfg.NUM_GPUS) == 0, \ 'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS) cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS effective_batch_size = args.iter_size * args.batch_size print('effective_batch_size = batch_size * iter_size = %d * %d' % (args.batch_size, args.iter_size)) print('Adaptive config changes:') print(' effective_batch_size: %d --> %d' % (original_batch_size, effective_batch_size)) print(' NUM_GPUS: %d --> %d' % (original_num_gpus, cfg.NUM_GPUS)) print(' IMS_PER_BATCH: %d --> %d' % (original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH)) print(' FG_THRESH: ', cfg.TRAIN.FG_THRESH) ### Adjust learning based on batch size change linearly # For iter_size > 1, gradients are `accumulated`, so lr is scaled based # on batch_size instead of effective_batch_size old_base_lr = cfg.SOLVER.BASE_LR cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size print('Adjust BASE_LR linearly according to batch_size change:\n' ' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR)) ### Adjust solver steps step_scale = original_batch_size / effective_batch_size old_solver_steps = cfg.SOLVER.STEPS old_max_iter = cfg.SOLVER.MAX_ITER cfg.SOLVER.STEPS = list( map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS)) cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5) cfg.SOLVER.VAL_ITER = int(cfg.SOLVER.VAL_ITER * step_scale + 0.5) cfg.TRAIN.SNAPSHOT_ITERS = int(cfg.TRAIN.SNAPSHOT_ITERS * step_scale + 0.5) print( 'Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n' ' SOLVER.STEPS: {} --> {}\n' ' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps, cfg.SOLVER.STEPS, old_max_iter, cfg.SOLVER.MAX_ITER)) # Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function # of `collect_and_distribute_fpn_rpn_proposals.py` # # post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5) if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN: cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch print( 'Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n' ' cfg.FPN.RPN_COLLECT_SCALE: {}'.format( cfg.FPN.RPN_COLLECT_SCALE)) if args.num_workers is not None: cfg.DATA_LOADER.NUM_THREADS = args.num_workers print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS) # ipdb.set_trace() ### Overwrite some solver settings from command line arguments if args.optimizer is not None: cfg.SOLVER.TYPE = args.optimizer if args.lr_decay_gamma is not None: cfg.SOLVER.GAMMA = args.lr_decay_gamma assert_and_infer_cfg() timers = defaultdict(Timer) ### Dataset ### timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) # Effective training sample size for one epoch train_size = roidb_size // args.batch_size * args.batch_size # ToDo: shuffle? batchSampler = BatchSampler(sampler=MinibatchSampler( ratio_list, ratio_index), batch_size=args.batch_size, drop_last=True) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_sampler=batchSampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) # dataiterator = iter(dataloader) ### Model ### from modeling.model_builder import Generalized_RCNN maskRCNN = Generalized_RCNN() if cfg.CUDA: maskRCNN.cuda() ### Optimizer ### bias_hoi_params = [] bias_hoi_param_names = [] bias_faster_params = [] bias_faster_param_names = [] nobias_hoi_params = [] nobias_hoi_param_names = [] nobias_faster_params = [] nobias_faster_param_names = [] # bias_params = [] # bias_param_names = [] # nonbias_params = [] # nonbias_param_names = [] #base_model = torch.load('Outputs/baseline/baseline_512_32_nogt_1o3/ckpt/model_step47999.pth') nograd_param_names = [] for key, value in maskRCNN.named_parameters(): #if key in base_model['model'].keys(): # value.requires_grad = False #print('the key xxx:', key) # Fix RPN module same as the paper # ToDo: or key.startswith('Box') # if 'affinity' not in key: # value.requires_grad = False print(key, value.size(), value.requires_grad) if value.requires_grad: if 'bias' in key: if 'HOI_Head' in key: bias_hoi_params.append(value) bias_hoi_param_names.append(key) else: bias_faster_params.append(value) bias_faster_param_names.append(key) else: if 'HOI_Head' in key: nobias_hoi_params.append(value) nobias_hoi_param_names.append(key) else: nobias_faster_params.append(value) nobias_faster_param_names.append(key) else: nograd_param_names.append(key) #del base_model #ipdb.set_trace() # Learning rate of 0 is a dummy value to be set properly at the start of training params = [ { 'params': nobias_hoi_params, 'lr': 0, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': nobias_faster_params, 'lr': 0 * cfg.SOLVER.FASTER_RCNN_WEIGHT, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': bias_hoi_params, 'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }, { 'params': bias_faster_params, 'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1) * cfg.SOLVER.FASTER_RCNN_WEIGHT, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }, ] if cfg.SOLVER.TYPE == "SGD": optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizer = torch.optim.Adam(params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) if args.krcnn_from_faster: net_utils.load_krcnn_from_faster(maskRCNN, checkpoint['model']) else: net_utils.load_ckpt(maskRCNN, checkpoint['model']) print('Original model loaded....') if args.resume: print('Resume, loaded step\n\n\n: ', checkpoint['step']) args.start_step = checkpoint['step'] + 1 if 'train_size' in checkpoint: # For backward compatibility if checkpoint['train_size'] != train_size: print( 'train_size value: %d different from the one in checkpoint: %d' % (train_size, checkpoint['train_size'])) # reorder the params in optimizer checkpoint's params_groups if needed # misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint) # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. optimizer.load_state_dict(checkpoint['optimizer']) # misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) del checkpoint torch.cuda.empty_cache() if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values) logging.info("loading Detectron weights %s", args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) lr = optimizer.param_groups[0][ 'lr'] # lr of non-bias parameters, for commmand line outputs. maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True) ### Training Setups ### args.run_name = misc_utils.get_run_name() + '_step' #output_dir = misc_utils.get_output_dir(args, args.run_name) output_dir = os.path.join('Outputs', args.expDir, args.expID) os.makedirs(output_dir, exist_ok=True) args.cfg_filename = os.path.basename(args.cfg_file) tblogger = None if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) if args.expID.__contains__('base'): os.environ['FABRICATOR'] = 'base' else: os.environ['FABRICATOR'] = 'fcl' print('log', os.environ['FABRICATOR']) ### Training Loop ### train_val(maskRCNN, args, optimizer, lr, dataloader, train_size, output_dir, tblogger)
def main(): """Main function""" args = parse_args() print('Called with args:') print(args) if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") if args.cuda or cfg.NUM_GPUS > 0: cfg.CUDA = True else: raise ValueError("Need Cuda device to run !") merge_cfg_from_file(args.cfg_file) # Some manual adjustment for the ApolloScape dataset parameters here cfg.OUTPUT_DIR = args.output_dir cfg.TRAIN.DATASETS = 'Car3D' cfg.MODEL.NUM_CLASSES = 8 if cfg.CAR_CLS.SIM_MAT_LOSS: cfg.MODEL.NUMBER_CARS = 79 else: # Loss is only cross entropy, hence, we detect only car categories in the training set. cfg.MODEL.NUMBER_CARS = 34 cfg.TRAIN.MIN_AREA = 196 # 14*14 cfg.TRAIN.USE_FLIPPED = False # Currently I don't know how to handle the flipped case cfg.TRAIN.IMS_PER_BATCH = 1 cfg.NUM_GPUS = torch.cuda.device_count() effective_batch_size = cfg.TRAIN.IMS_PER_BATCH * cfg.NUM_GPUS * args.iter_size ### Adaptively adjust some configs ### original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH original_num_gpus = cfg.NUM_GPUS if args.batch_size is None: args.batch_size = original_batch_size assert (args.batch_size % cfg.NUM_GPUS) == 0, 'batch_size: %d, NUM_GPUS: %d' % ( args.batch_size, cfg.NUM_GPUS) print('effective_batch_size = batch_size * iter_size = %d * %d' % (args.batch_size, args.iter_size)) print('Adaptive config changes:') print(' effective_batch_size: %d --> %d' % (original_batch_size, effective_batch_size)) print(' NUM_GPUS: %d --> %d' % (original_num_gpus, cfg.NUM_GPUS)) print(' IMS_PER_BATCH: %d --> %d' % (original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH)) ### Adjust learning based on batch size change linearly # For iter_size > 1, gradients are `accumulated`, so lr is scaled based # on batch_size instead of effective_batch_size old_base_lr = cfg.SOLVER.BASE_LR cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size print( 'Adjust BASE_LR linearly according to batch_size change:\n BASE_LR: {} --> {}' .format(old_base_lr, cfg.SOLVER.BASE_LR)) ### Adjust solver steps step_scale = original_batch_size / effective_batch_size old_solver_steps = cfg.SOLVER.STEPS old_max_iter = cfg.SOLVER.MAX_ITER cfg.SOLVER.STEPS = list( map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS)) cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5) print( 'Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n' ' SOLVER.STEPS: {} --> {}\n' ' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps, cfg.SOLVER.STEPS, old_max_iter, cfg.SOLVER.MAX_ITER)) # Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function # of `collect_and_distribute_fpn_rpn_proposals.py` # # post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5) if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN: cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch print( 'Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n' ' cfg.FPN.RPN_COLLECT_SCALE: {}'.format( cfg.FPN.RPN_COLLECT_SCALE)) if args.num_workers is not None: cfg.DATA_LOADER.NUM_THREADS = args.num_workers print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS) ### Overwrite some solver settings from command line arguments if args.optimizer is not None: cfg.SOLVER.TYPE = args.optimizer if args.lr is not None: cfg.SOLVER.BASE_LR = args.lr if args.lr_decay_gamma is not None: cfg.SOLVER.GAMMA = args.lr_decay_gamma assert_and_infer_cfg() timers = defaultdict(Timer) ### Dataset ### timers['roidb'].tic() if cfg.MODEL.LOSS_3D_2D_ON: roidb, ratio_list, ratio_index, ds = combined_roidb_for_training( cfg.TRAIN.DATASETS, args.dataset_dir) else: roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, args.dataset_dir) timers['roidb'].toc() roidb_size = len(roidb) logger.info('{:d} roidb entries'.format(roidb_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) # Effective training sample size for one epoch train_size = roidb_size // args.batch_size * args.batch_size sampler = MinibatchSampler(ratio_list, ratio_index) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True, valid_keys=[ 'has_visible_keypoints', 'boxes', 'seg_areas', 'gt_classes', 'gt_overlaps', 'box_to_gt_ind_map', 'is_crowd', 'car_cat_classes', 'poses', 'quaternions', 'im_info' ]) dataloader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, drop_last=True, sampler=sampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) dataiterator = iter(dataloader) ### Model ### if cfg.MODEL.LOSS_3D_2D_ON: maskRCNN = Generalized_RCNN(ds.Car3D) else: maskRCNN = Generalized_RCNN() if cfg.CUDA: maskRCNN.cuda() ### Optimizer ### bias_params = [] bias_param_names = [] nonbias_params = [] nonbias_param_names = [] for key, value in dict(maskRCNN.named_parameters()).items(): if value.requires_grad: if 'bias' in key: bias_params.append(value) bias_param_names.append(key) else: nonbias_params.append(value) nonbias_param_names.append(key) # Learning rate of 0 is a dummy value to be set properly at the start of training params = [{ 'params': nonbias_params, 'lr': 0, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': bias_params, 'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }] if cfg.SOLVER.TYPE == "SGD": optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizer = torch.optim.Adam(params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model'], ignore_list=args.ckpt_ignore_head) if args.resume: args.start_step = checkpoint['step'] + 1 if 'train_size' in checkpoint: # For backward compatibility if checkpoint['train_size'] != train_size: print( 'train_size value: %d different from the one in checkpoint: %d' % (train_size, checkpoint['train_size'])) # reorder the params in optimizer checkpoint's params_groups if needed # misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint) # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. # optimizer.load_state_dict(checkpoint['optimizer']) misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) del checkpoint torch.cuda.empty_cache() if args.load_detectron: # TODO resume for detectron weights (load sgd momentum values) logging.info("loading Detectron weights %s", args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) lr = optimizer.param_groups[0][ 'lr'] # lr of non-bias parameters, for commmand line outputs. maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True) ### Training Setups ### args.run_name = misc_utils.get_run_name() + '_step' # output_dir = os.path.join('/media/SSD_1TB/zzy/ApolloScape/ECCV2018_apollo/train', args.run_name) output_dir = misc_utils.get_output_dir(args, args.run_name) args.cfg_filename = os.path.basename(args.cfg_file) if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) ### Training Loop ### maskRCNN.train() CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS) # Set index for decay steps decay_steps_ind = None for i in range(1, len(cfg.SOLVER.STEPS)): if cfg.SOLVER.STEPS[i] >= args.start_step: decay_steps_ind = i break if decay_steps_ind is None: decay_steps_ind = len(cfg.SOLVER.STEPS) training_stats = TrainingStats( args, args.disp_interval, tblogger if args.use_tfboard and not args.no_save else None) # warmup_factor_trans = 1.0 try: logger.info('Training starts !') step = args.start_step for step in range(args.start_step, cfg.SOLVER.MAX_ITER): # Warm up if step < cfg.SOLVER.WARM_UP_ITERS: method = cfg.SOLVER.WARM_UP_METHOD if method == 'constant': warmup_factor = cfg.SOLVER.WARM_UP_FACTOR elif method == 'linear': alpha = step / cfg.SOLVER.WARM_UP_ITERS warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 - alpha) + alpha # warmup_factor_trans = cfg.SOLVER.WARM_UP_FACTOR_TRANS * (1 - alpha) + alpha # warmup_factor_trans *= cfg.TRANS_HEAD.LOSS_BETA warmup_factor_trans = 1.0 else: raise KeyError( 'Unknown SOLVER.WARM_UP_METHOD: {}'.format(method)) lr_new = cfg.SOLVER.BASE_LR * warmup_factor net_utils.update_learning_rate(optimizer, lr, lr_new) lr = optimizer.param_groups[0]['lr'] assert lr == lr_new elif step == cfg.SOLVER.WARM_UP_ITERS: net_utils.update_learning_rate(optimizer, lr, cfg.SOLVER.BASE_LR) lr = optimizer.param_groups[0]['lr'] assert lr == cfg.SOLVER.BASE_LR # Learning rate decay if decay_steps_ind < len( cfg.SOLVER.STEPS ) and step == cfg.SOLVER.STEPS[decay_steps_ind]: logger.info('Decay the learning on step %d', step) lr_new = lr * cfg.SOLVER.GAMMA net_utils.update_learning_rate(optimizer, lr, lr_new) lr = optimizer.param_groups[0]['lr'] assert lr == lr_new decay_steps_ind += 1 training_stats.IterTic() optimizer.zero_grad() for inner_iter in range(args.iter_size): try: input_data = next(dataiterator) except StopIteration: dataiterator = iter(dataloader) input_data = next(dataiterator) for key in input_data: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list(map(Variable, input_data[key])) net_outputs = maskRCNN(**input_data) net_outputs['losses'][ 'loss_car_cls'] *= cfg.CAR_CLS.CAR_CLS_LOSS_BETA net_outputs['losses']['loss_rot'] *= cfg.CAR_CLS.ROT_LOSS_BETA if cfg.MODEL.TRANS_HEAD_ON: net_outputs['losses'][ 'loss_trans'] *= cfg.TRANS_HEAD.TRANS_LOSS_BETA training_stats.UpdateIterStats_car_3d(net_outputs) # start training # loss_car_cls: 2.233790, loss_rot: 0.296853, loss_trans: ~100 loss = net_outputs['losses']['loss_car_cls'] + net_outputs[ 'losses']['loss_rot'] if cfg.MODEL.TRANS_HEAD_ON: loss += net_outputs['losses']['loss_trans'] if cfg.MODEL.LOSS_3D_2D_ON: loss += net_outputs['losses']['UV_projection_loss'] if not cfg.TRAIN.FREEZE_CONV_BODY and not cfg.TRAIN.FREEZE_RPN and not cfg.TRAIN.FREEZE_FPN: loss += net_outputs['total_loss_conv'] loss.backward() optimizer.step() training_stats.IterToc() training_stats.LogIterStats(step, lr, warmup_factor_trans) if (step + 1) % CHECKPOINT_PERIOD == 0: save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) # ---- Training ends ---- # Save last checkpoint save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) except (RuntimeError, KeyboardInterrupt): del dataiterator logger.info('Save ckpt on exception ...') save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer) logger.info('Save ckpt done.') stack_trace = traceback.format_exc() print(stack_trace) finally: if args.use_tfboard and not args.no_save: tblogger.close()