def load_model(self, weight_file, use_half, device_id): # import pytorch # self._log.info("load_model does not use device, or use_half in MRCNN") # print("Device in load model: ", device) try: import torch except: raise RuntimeError("could not load pytorch!") # import model try: from modeling.model_builder import Generalized_RCNN as ubMRCNN # from ubmrcnn import ubMRCNN except: raise RuntimeError( "could not load ubMRCNN model. did you remember" + " to setup everything?") # if "cuda" not in device and "cpu" not in device: # raise ValueError("invalid device name [{}]. Must str with name \ # \"cpu\" or \"cuda:X\" where X=device number") self._log = logging.getLogger(self.idname()) # self.device = torch.device(device) self.model = Generalized_RCNN() #checkpoint = torch.load(weight_file, map_location=lambda storage, loc: storage) locations = {} for x in range(6): locations["cuda:%d" % (x)] = "cpu" checkpoint = torch.load(weight_file, map_location=locations) # self.device = device_id # if device_id==None: # self.device = torch.device("cpu") # else: # self.device = torch.device("cuda:%d"%(device_id)) net_utils.load_ckpt(self.model, checkpoint['model']) # self.model = mynn.DataParallel(self.model, cpu_keywords=['im_info', 'roidb'], # minibatch=True, device_ids=[0], # output_device=0) # only support single GPU self.model = mynn.DataSingular(self.model, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_id=[cfg.MODEL.DEVICE]) self.model.eval()
def initialize_model_from_cfg(args, gpu_id=0): """Initialize a model from the global cfg. Loads test-time weights and set to evaluation mode. """ model = Generalized_RCNN() model.eval() if args.cuda: model.cuda() if args.load_ckpt: load_name = args.load_ckpt logger.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(model, checkpoint['model']) # model.load_state_dict(checkpoint['model']) if args.load_detectron: logger.info("loading detectron weights %s", args.load_detectron) load_detectron_weight(model, args.load_detectron) model = mynn.DataParallel(model, cpu_keywords=['im_info', 'roidb'], minibatch=True) return model
def __init__(self): args = parse_args() # Historical code, because we have trained on the Baidu Apoloscape dataset. dataset = WAD_CVPR2018(args.dataset_dir) cfg.MODEL.NUM_CLASSES = len( dataset.eval_class) + 1 # with a background class print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(osp.join(this_dir, args.cfg_file)) if args.nms_soft: cfg.TEST.SOFT_NMS.ENABLED = True else: cfg.TEST.NMS = args.nms cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() maskRCNN.cuda() if args.load_ckpt: load_name = osp.join(this_dir, args.load_ckpt) print("loading checkpoint %s" % (load_name)) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) maskRCNN.eval() self.model = maskRCNN self.dataset = dataset
def parse_dataset(path_to_dataset): ''' Parse every image contained in 'path_to_dataset' (a path to the training or testing set), classify each image and save in a csv file the resulting assignment dataset_result: dict structure returned by the function. It contains the label of each image ''' llist = listdir(path_to_dataset) dataset_result = {} maskRCNN = Generalized_RCNN() maskRCNN.cuda() maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() for name in tqdm(llist): dn = join(path_to_dataset, name) timers = defaultdict(Timer) dataset_result[dn] = localize_obj_in_image(dn, maskRCNN, timers) return dataset_result
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) dataset = WAD_CVPR2018(args.dataset_dir) cfg.MODEL.NUM_CLASSES = len( dataset.eval_class) + 1 # with a background class print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("loading checkpoint %s" % (load_name)) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() imglist = misc_utils.get_imagelist_from_dir(dataset.test_image_dir) num_images = len(imglist) output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]), 'Images') if not os.path.exists(output_dir): os.makedirs(output_dir) for i in tqdm(xrange(num_images)): im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) args.current_im_name = im_name cls_boxes, cls_segms, prediction_row = im_detect_all(args, maskRCNN, im, dataset, timers=timers) thefile = open(os.path.join(output_dir, im_name + '.txt'), 'w') for item in prediction_row: thefile.write("%s\n" % item)
def GetRCNNModel(): """ RCNN model factory """ if cfg.CONTEXT_ROI_POOLING.ENABLED: maskRCNN = Context_RCNN() elif cfg.CONTRASTED_CONTEXT.ENABLED: maskRCNN = CC_RCNN() elif cfg.LESION.USE_3DCE: maskRCNN = CE3D_RCNN() elif cfg.LESION.MULTI_MODALITY: maskRCNN = MULTI_MODALITY_RCNN() elif cfg.LESION.MM_TEST: maskRCNN = MM_TEST_RCNN() else: maskRCNN = Generalized_RCNN() return maskRCNN
def __init__(self): cfg_from_file("maskrcnn/configs/baselines/e2e_mask_rcnn_X-152-32x8d-FPN-IN5k_1.44x.yaml") cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() maskRCNN.cuda() pretrained_path = "maskrcnn/data/X-152-32x8d-FPN-IN5k.pkl" print("loading detectron weights %s" % pretrained_path) load_detectron_weight(maskRCNN, pretrained_path) maskRCNN.eval() # Note this step self.maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset == "pascal_parts_heads": dataset = datasets.get_head_dataset() cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "scuthead_a": dataset = datasets.get_head_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) for i in xrange(num_images): im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) outputfile = os.path.join(args.output_dir, im_name) head_boxes = cls_boxes[1] print('img :', i, ' num_heads :', len(head_boxes), ' img_path :', imglist[i]) np.save(outputfile, head_boxes)
# ----------------------------------------------------------------------------------- # Model setup # ----------------------------------------------------------------------------------- cfg.TEST.SCALE = 800 cfg.TEST.MAX_SIZE = 1333 cfg.MODEL.NUM_CLASSES = 2 cfg.TEST.NMS = NMS_THRESH print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False assert_and_infer_cfg() net = Generalized_RCNN() if args.cuda: net.cuda() if args.load_ckpt: load_name = args.load_ckpt print("loading checkpoint %s" % (load_name)) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(net, checkpoint['model']) if args.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(net, args.load_detectron)
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 elif args.dataset == "rsna2018": dataset = datasets.get_rsna_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) submission_dict = {} vis_dir = os.path.join(args.output_dir, 'vis') for i in xrange(num_images): print('img', i) predictionstring = '' im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) boxes = [] # vis_utils.vis_one_image( # im[:, :, ::-1], # BGR -> RGB for visualization # im_name, # vis_dir, # cls_boxes, # cls_segms, # cls_keyps, # dataset=dataset, # box_alpha=0.3, # show_class=True, # thresh=args.submit_threshold, # kp_thresh=2, # ext='png' # ) box_list = [b for b in cls_boxes if len(b) > 0] if len(box_list) > 0: boxes = np.concatenate(box_list) for box in boxes: bbox = box[:4] score = box[-1] if score < args.submit_threshold: continue x = bbox[0] y = bbox[1] height = bbox[3] - bbox[1] width = bbox[2] - bbox[0] predictionstring += str(score) + ' ' + str(x) + ' ' + str( y) + ' ' + str(width) + ' ' + str(height) + ' ' filename = imglist[i].split('.')[0] filename = filename.split('/')[-1] submission_dict[filename] = predictionstring assert len(submission_dict) <= num_images sub = pd.DataFrame.from_dict(submission_dict, orient='index') sub.index.name = 'patientId' sub.columns = ['PredictionString'] sub.to_csv(os.path.join(args.output_dir, "submission.csv"))
def test_net_on_dataset(args): dataset = WAD_CVPR2018(args.dataset_dir) cfg.MODEL.NUM_CLASSES = len(dataset.eval_class) + 1 # with a background class print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(osp.join(this_dir,args.cfg_file)) if args.nms_soft: cfg.TEST.SOFT_NMS.ENABLED = True else: cfg.TEST.NMS = args.nms cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() maskRCNN.cuda() if args.load_ckpt: load_name = osp.join(this_dir, args.load_ckpt) print("loading checkpoint %s" % (load_name)) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.nms_soft: output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]), 'Images_' + str(cfg.TEST.SCALE)+'_SOFT_NMS') elif args.nms: output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]), 'Images_' + str(cfg.TEST.SCALE)+'_NMS_%.2f'%args.nms) else: output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]), 'Images_' + str(cfg.TEST.SCALE)) if cfg.TEST.BBOX_AUG.ENABLED: output_dir += '_TEST_AUG' #if args.cls_boxes_confident_threshold < 0.5: output_dir += '_cls_boxes_confident_threshold_%.1f' % args.cls_boxes_confident_threshold if not os.path.exists(output_dir): os.makedirs(output_dir) args.output_img_dir = os.path.join(output_dir, 'Image_Masks') if not os.path.exists(args.output_img_dir): os.makedirs(args.output_img_dir) output_list_dir = '/home/stevenwudi/PycharmProjects/Udacity/CarND-Vehicle-Detection/output_images/car_detection_maskrcnn' if not os.path.exists(output_list_dir): os.makedirs(output_list_dir) output_vis_dir = output_list_dir imglist = ['/home/stevenwudi/PycharmProjects/Udacity/CarND-Vehicle-Detection/test_images/test3.jpg', '/home/stevenwudi/PycharmProjects/Udacity/CarND-Vehicle-Detection/test_images/test4.jpg', '/home/stevenwudi/PycharmProjects/Udacity/CarND-Vehicle-Detection/test_images/test5.jpg'] if not args.range: start = 0 end = len(imglist) else: start = args.range[0] end = args.range[1] for i in tqdm(xrange(start, end)): im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) args.current_im_name = im_name cls_boxes, cls_segms, prediction_row = im_detect_all(args, maskRCNN, im, dataset, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) print(im_name) if args.vis: vis_utils.vis_one_image_cvpr2018_wad( im[:, :, ::-1], # BGR -> RGB for visualization im_name, output_vis_dir, cls_boxes, cls_segms, None, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.99, kp_thresh=2 ) thefile = open(os.path.join(output_list_dir, im_name + '.txt'), 'w') for item in prediction_row: thefile.write("%s\n" % item)
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()
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) for i in xrange(num_images): print('img', i) im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes, cls_segms, cls_keyps, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2) if args.merge_pdfs and num_images > 1: merge_out_path = '{}/results.pdf'.format(args.output_dir) if os.path.exists(merge_out_path): os.remove(merge_out_path) command = "pdfunite {}/*.pdf {}".format(args.output_dir, merge_out_path) subprocess.call(command, shell=True)
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()
p_tensor.copy_(torch.Tensor(src_blobs[d_name])) def resnet_weights_name_pattern(): pattern = re.compile(r"conv1_w|conv1_gn_[sb]|res_conv1_.+|res\d_\d_.+") return pattern if __name__ == '__main__': """Testing""" from pprint import pprint import sys sys.path.insert(0, '..') from modeling.model_builder import Generalized_RCNN from core.config import cfg, cfg_from_file cfg.MODEL.NUM_CLASSES = 81 cfg_from_file('../../cfgs/res50_mask.yml') net = Generalized_RCNN() # pprint(list(net.state_dict().keys()), width=1) mapping, orphans = net.detectron_weight_mapping state_dict = net.state_dict() for k in mapping.keys(): assert k in state_dict, '%s' % k rest = set(state_dict.keys()) - set(mapping.keys()) assert len(rest) == 0
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 elif args.dataset.startswith("gangjin"): dataset = datasets.get_gangjin_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) img_ids = [] rects = [] for i in range(num_images): print('img', i) im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) boxes, segms, keypoints, classes = vis_utils.convert_from_cls_format( cls_boxes, cls_segms, cls_keyps) if boxes is not None: for j in range(len(boxes)): # print(boxes[j][-1]) if float(boxes[j][-1]) < 0.99: # 阀值 continue xmin = float(boxes[j, 0]) xmax = float(boxes[j, 2]) ymin = float(boxes[j, 1]) ymax = float(boxes[j, 3]) img_ids.append(os.path.basename(imglist[i])) rects.append( str(xmin) + " " + str(ymin) + " " + str(xmax) + " " + str(ymax)) # im_name, _ = os.path.splitext(os.path.basename(imglist[i])) # vis_utils.vis_one_image( # im[:, :, ::-1], # BGR -> RGB for visualization # im_name, # args.output_dir, # cls_boxes, # cls_segms, # cls_keyps, # dataset=dataset, # box_alpha=0.3, # show_class=False, # thresh=0.99, # kp_thresh=2, # ext="jpg" # ) result_dict = {"ID": img_ids, "rects": rects} import pandas as pd result = pd.DataFrame.from_dict(result_dict) result.to_csv('submit/submit1.csv', header=None, index=False)
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) prefix_path = args.output_dir + '_results' if os.path.exists(prefix_path): shutil.rmtree(prefix_path) os.mkdir(prefix_path) else: os.mkdir(prefix_path) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) for i in tqdm(range(num_images)): im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) boxes, _, _, classes = convert_from_cls_format(cls_boxes, cls_segms, cls_keyps) if classes == []: continue voc_boxes = np.zeros_like(boxes) voc_boxes[:, 0:1] = boxes[:, 4:5] voc_boxes[:, 1:3] = boxes[:, 0:2] + 1 voc_boxes[:, 3:5] = boxes[:, 2:4] + 1 for instance_idx, cls_idx in enumerate(classes): cls_name = dataset.classes[cls_idx] if cls_name == 'motorcycle': cls_name = 'motorbike' f = open(os.path.join(prefix_path, cls_name + ".txt"), "a+") f.write("%s " % im_name) for item in voc_boxes[instance_idx]: f.write("%f " % item) f.write("\n") f.close()
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) prefix_path = args.output_dir os.makedirs(prefix_path, exist_ok=True) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) writen_results = [] # validate demo_im = cv2.imread(imglist[0]) print(np.shape(demo_im)) h, w, _ = np.shape(demo_im) #print(h) #print(args.height) assert h == args.height assert w == args.width h_scale = 720 / args.height w_scale = 1280 / args.width for i in tqdm(range(num_images)): im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) # boxs = [[x1, y1, x2, y2, cls], ...] boxes, _, _, classes = convert_from_cls_format(cls_boxes, cls_segms, cls_keyps) if boxes is None: continue # scale boxes[:, 0] = boxes[:, 0] * w_scale boxes[:, 2] = boxes[:, 2] * w_scale boxes[:, 1] = boxes[:, 1] * h_scale boxes[:, 3] = boxes[:, 3] * h_scale if classes == []: continue for instance_idx, cls_idx in enumerate(classes): cls_name = dataset.classes[cls_idx] if cls_name == 'motorcycle': cls_name = 'motor' elif cls_name == 'stop sign': cls_name = 'traffic sign' elif cls_name == 'bicycle': cls_name = 'bike' if cls_name not in bdd_category: continue writen_results.append({ "name": imglist[i].split('/')[-1], "timestamp": 1000, "category": cls_name, "bbox": boxes[instance_idx, :4], "score": boxes[instance_idx, -1] }) with open(os.path.join(prefix_path, args.name + '.json'), 'w') as outputfile: json.dump(writen_results, outputfile, cls=MyEncoder)
from __future__ import absolute_import
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset == "miotcd": dataset = datasets.get_miotcd_dataset() cfg.MODEL.NUM_CLASSES = 12 elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 elif args.dataset.startswith("bogota"): dataset = datasets.get_bogota_dataset() cfg.MODEL.NUM_CLASSES = 12 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) dataset_result = {} for i in xrange(num_images): print('img', i) im = cv2.imread(imglist[i]) assert im is not None try: timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) boxes_, segme_ , keyps_ ,clasies= convert_from_cls_format(cls_boxes,cls_segms,cls_keyps) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) dataset_result[im_name] = localize_obj_in_image(im_name,boxes_,clasies) except(e): import pdb pdb.set_trace() np.save('dictionary_answer.npy', dataset_result) tmp = args.image_dir.split('/') txt = str(tmp[-2:])+'_'+str(tmp[-1:])+'.csv' save_localization_result(dataset_result,txt)
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_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 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() + '_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()
import glob import subprocess from PIL import Image from torchvision import transforms import torch.nn.functional as F # load config dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) cfg_from_file('configs/baselines/e2e_mask_rcnn_R-50-C4_1x.yaml') cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() # load model maskRCNN = Generalized_RCNN() checkpoint = torch.load( '/home/work/liupeng11/code/Detectron.pytorch/models/e2e_mask_rcnn_R-50-C4_1x.pth', map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) maskRCNN.eval() maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # load image img_path = "/home/work/liupeng11/code/Detectron.pytorch/demo/sample_images/img1.jpg" im = cv2.imread(img_path) # detect bouding boxes and segments
def main(): """main function""" os.environ['CUDA_VISIBLE_DEVICES'] = '1' if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() params = list(maskRCNN.parameters()) k = 0 for i in params: l = 1 for j in i.size(): l *= j k = k + l print('zonghe:' + str(k)) if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) for i in xrange(num_images): print('img', i) im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) start = time.time() cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) class_result_boxes = [] for index, class_boxes in enumerate(cls_boxes): if len(class_boxes) != 0: class_boxes = class_boxes.tolist() results_oneclass = threeD_detect(imglist[i], class_boxes, index) class_result_boxes.append(results_oneclass) save_image = im color_class = { 'Car': [0, 255, 255], 'Cyclist': [255, 0, 0], 'Pedestrian': [0, 0, 255] } for result_boxes in class_result_boxes: for box in result_boxes: cv2.rectangle(save_image, (box[0], box[1]), (box[2], box[3]), color_class[box[-1]], 2) height = round(box[-2][0], 2) width = round(box[-2][1], 2) length = round(box[-2][2], 2) threeD_info = str(height) + ' ' + str(width) + ' ' + str( length) cv2.putText(save_image, threeD_info, (box[0], box[1] - 20), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 2) _, imagename = os.path.split(imglist[i]) imagename2 = imagename.split('.')[0] cv2.imwrite('../output1/%s.png' % imagename2, save_image) end = time.time() print(end - start) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes, cls_segms, cls_keyps, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2) if args.merge_pdfs and num_images > 1: merge_out_path = '{}/results.pdf'.format(args.output_dir) if os.path.exists(merge_out_path): os.remove(merge_out_path) command = "pdfunite {}/*.pdf {}".format(args.output_dir, merge_out_path) subprocess.call(command, shell=True)
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) for i in xrange(num_images): print('img', i) im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes, cls_segms, cls_keyps, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2 ) if args.merge_pdfs and num_images > 1: merge_out_path = '{}/results.pdf'.format(args.output_dir) if os.path.exists(merge_out_path): os.remove(merge_out_path) command = "pdfunite {}/*.pdf {}".format(args.output_dir, merge_out_path) subprocess.call(command, shell=True)
def test_net_on_dataset(args): dataset = WAD_CVPR2018(args.dataset_dir) cfg.MODEL.NUM_CLASSES = len( dataset.eval_class) + 1 # with a background class print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.nms_soft: cfg.TEST.SOFT_NMS.ENABLED = True else: cfg.TEST.NMS = args.nms cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("loading checkpoint %s" % (load_name)) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() imglist_all = misc_utils.get_imagelist_from_dir(dataset.test_image_dir) if args.nms_soft: output_dir = os.path.join( ('/').join(args.load_ckpt.split('/')[:-2]), 'Images_' + str(cfg.TEST.SCALE) + '_SOFT_NMS') elif args.nms: output_dir = os.path.join( ('/').join(args.load_ckpt.split('/')[:-2]), 'Images_' + str(cfg.TEST.SCALE) + '_NMS_%.2f' % args.nms) else: output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]), 'Images_' + str(cfg.TEST.SCALE)) if cfg.TEST.BBOX_AUG.ENABLED: output_dir += '_TEST_AUG' #if args.cls_boxes_confident_threshold < 0.5: output_dir += '_cls_boxes_confident_threshold_%.1f' % args.cls_boxes_confident_threshold if not os.path.exists(output_dir): os.makedirs(output_dir) output_vis_dir = os.path.join(output_dir, 'Image_Vis') if not os.path.exists(output_vis_dir): os.makedirs(output_vis_dir) args.output_img_dir = os.path.join(output_dir, 'Image_Masks') if not os.path.exists(args.output_img_dir): os.makedirs(args.output_img_dir) output_list_dir = os.path.join(output_dir, 'List_Masks') if not os.path.exists(output_list_dir): os.makedirs(output_list_dir) # A break point # img_produced = os.listdir(output_list_dir) # imglist = [x for x in imglist_all if x.split('/')[-1][:-4]+'.txt' not in img_produced] imglist = imglist_all if False: args.vis = True imglist = [ os.path.join(args.dataset_dir, 'test', '4f38c1d630209cfb777a3dcbb613ba56.jpg'), os.path.join(args.dataset_dir, 'test', '3b08ff3b969200982d0283aedbcaae20.jpg') ] if not args.range: start = 0 end = len(imglist) else: start = args.range[0] end = args.range[1] for i in tqdm(xrange(start, end)): im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) args.current_im_name = im_name cls_boxes, cls_segms, prediction_row = im_detect_all(args, maskRCNN, im, dataset, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) print(im_name) if args.vis: vis_utils.vis_one_image_cvpr2018_wad( im[:, :, ::-1], # BGR -> RGB for visualization im_name, output_vis_dir, cls_boxes, cls_segms, None, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.5, kp_thresh=2) thefile = open(os.path.join(output_list_dir, im_name + '.txt'), 'w') for item in prediction_row: thefile.write("%s\n" % item)
def test_net_on_dataset_multigpu(args): dataset = WAD_CVPR2018(args.dataset_dir) cfg.MODEL.NUM_CLASSES = len( dataset.eval_class) + 1 # with a background class print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights if args.nms_soft: cfg.TEST.SOFT_NMS.ENABLED = True else: cfg.TEST.NMS = args.nms assert_and_infer_cfg() maskRCNN = Generalized_RCNN() maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("loading checkpoint %s" % (load_name)) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() #imglist = misc_utils.get_imagelist_from_dir(dataset.test_image_dir) img_produced = os.listdir(args.output_list_dir) imglist_all = misc_utils.get_imagelist_from_dir(dataset.test_image_dir) imglist = [ x for x in imglist_all if x.split('/')[-1][:-4] + '.txt' not in img_produced ] #for i in tqdm(xrange(args.range[0], args.range[1])): start_idx = len(imglist_all) - args.range[0] end_idx = len(imglist_all) - args.range[0] + args.range[1] for i in tqdm(xrange(start_idx, end_idx)): im = cv2.imread(imglist_all[i]) assert im is not None timers = defaultdict(Timer) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) args.current_im_name = im_name cls_boxes, cls_segms, prediction_row = im_detect_all(args, maskRCNN, im, dataset, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) print(im_name) if args.vis: vis_utils.vis_one_image_cvpr2018_wad( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_vis_dir, cls_boxes, cls_segms, None, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.5, kp_thresh=2) thefile = open(os.path.join(args.output_list_dir, im_name + '.txt'), 'w') for item in prediction_row: thefile.write("%s\n" % item)
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) video_path = args.video video_save_dir = args.video_save_dir if not os.path.exists(video_save_dir): os.mkdir(video_save_dir) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() capture = cv2.VideoCapture(video_path) frame_count = 0 batch_size = 1 frame_list = [] while True: ret, frame = capture.read() # Bail out when the video file ends if not ret: break # Save each frame of the video to a list frame_count += 1 frame_list.append(frame) if len(frame_list) == batch_size: for i, frame in enumerate(frame_list): timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, frame, timers=timers) name = '{0}'.format(frame_count + i - batch_size) print('\n' + name) vis_utils.vis_one_image( frame[:, :, ::-1], # BGR -> RGB for visualization name, video_save_dir, cls_boxes, cls_segms, cls_keyps, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.8, kp_thresh=2, ext='jpg') # Clear the frames array to start the next batch frame_list = [] images = list(glob.iglob(os.path.join(video_save_dir, '*.jpg'))) # Sort the images by integer index images = sorted(images, key=lambda x: float(os.path.split(x)[1][:-3])) outvid = os.path.join(video_save_dir, "out.mp4") make_video(outvid, images, fps=30) extract_audio = 'ffmpeg -i %s -vn -acodec copy %s/out.aac' % ( video_path, video_save_dir) subprocess.call(extract_audio, shell=True) merge_audio_video = "ffmpeg -i %s/out.aac -i %s/out.mp4 -codec copy -shortest %s/final.mp4" % ( video_save_dir, video_save_dir, video_save_dir) subprocess.call(merge_audio_video, shell=True)
p_tensor.copy_(torch.Tensor(src_blobs[d_name])) def resnet_weights_name_pattern(): pattern = re.compile(r"conv1_w|conv1_gn_[sb]|res_conv1_.+|res\d+_\d+_.+") return pattern if __name__ == '__main__': """Testing""" from pprint import pprint import sys sys.path.insert(0, '..') from modeling.model_builder import Generalized_RCNN from core.config import cfg, cfg_from_file cfg.MODEL.NUM_CLASSES = 81 cfg_from_file('../../cfgs/res50_mask.yml') net = Generalized_RCNN() # pprint(list(net.state_dict().keys()), width=1) mapping, orphans = net.detectron_weight_mapping state_dict = net.state_dict() for k in mapping.keys(): assert k in state_dict, '%s' % k rest = set(state_dict.keys()) - set(mapping.keys()) assert len(rest) == 0
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) dataset = datasets.get_hospital_dataset() cfg.MODEL.NUM_CLASSES = 20 # with bg num_class = cfg.MODEL.NUM_CLASSES sents = dataset.sents th_cls = dataset.th_cls cls2eng = dataset.cls2eng eng2type = dataset.eng2type print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("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.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) for i in xrange(num_images): # for each image print('img', i) im = cv2.imread(imglist[i]) assert im is not None # segmentation # d = segment(im) # pdb.set_trace() timers = defaultdict(Timer) # detection cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) # first we collect boxes from all classes dets_total = np.empty([0, 6], dtype=np.float32) for cls in range(1, num_class): # for each cls dets = cls_boxes[cls] if dets.shape[0] == 0: continue dets_extend = np.pad( dets, ((0, 0), (0, 1)), # add 0 rows above, below and left, but 1 row right mode='constant', constant_values=cls) # append cls to dets dets_total = np.vstack((dets_total, dets_extend)) # then use a loose NMS to make each region has only one symptom keep = box_utils.nms(dets_total, 0.7) nms_dets = dets_total[keep, :] # iterate through remained boxes report, healthy = '', True have_sym_of_cls = [False for _ in range(num_class)] n = nms_dets.shape[0] final_results = [] # return to the web for idx in range(n): # for each region th, cls = nms_dets[idx, -2], int(nms_dets[idx, -1]) if th > th_cls[cls]: # diagnosed to have the sym report += sents[cls][1] have_sym_of_cls[cls] = True healthy = False ename = cls2eng[int(cls)] _type = eng2type[ename] final_results.append({ 'name': ename, 'type': _type, 'box': list(nms_dets[idx, 0:4]) }) for cls in range(1, num_class): # for each cls if not have_sym_of_cls[cls]: # if have no sym of this cls report += sents[cls][0] if healthy: report = sents[0][0] print(report) pdb.set_trace() # healthy = True # flag indicating healthy or not # for cls in range(1, num_class): # for each cls # dets = cls_boxes[cls] # if dets.shape[0] == 0: # report += sents[cls][0] # continue # n = dets.shape[0] # flag = False # indicates if the sym exists # for k in range(n): # for each region # if dets[k, -1] > th_cls[cls]: # above threshold for this cls, means have this cls of symptom # report += sents[cls][1] # flag = True # healthy = False # if not flag: # don't have this symptom # report += sents[cls][0] # # if healthy: # use the report for healthy people # report = sents[0][0] im_name, _ = os.path.splitext(os.path.basename(imglist[i])) # vis_utils.vis_one_image( # im[:, :, ::-1], # BGR -> RGB for visualization # im_name, # args.output_dir, # cls_boxes, # cls_segms, # cls_keyps, # dataset=dataset, # box_alpha=0.3, # show_class=True, # thresh=0.05, # kp_thresh=2 # ) if args.merge_pdfs and num_images > 1: merge_out_path = '{}/results.pdf'.format(args.output_dir) if os.path.exists(merge_out_path): os.remove(merge_out_path) command = "pdfunite {}/*.pdf {}".format(args.output_dir, merge_out_path) subprocess.call(command, shell=True)
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""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 elif args.dataset.startswith("keypoints_carfusion"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) assert bool(args.load_ckpt_car) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' print('load cfg from file: {}'.format(args.cfg_file_person)) cfg_from_file(args.cfg_file_person) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN_person = Generalized_RCNN() #print('load cfg from file: {}'.format(args.cfg_file_person)) #cfg_from_file(args.cfg_file_person) #assert_and_infer_cfg() #maskRCNN_person = Generalized_RCNN() if args.visualize: save_image = True if args.cuda: maskRCNN_person.cuda() if args.load_ckpt_person: load_name = args.load_ckpt_person print("loading checkpoint for person %s" % (load_name)) checkpoint_person = torch.load( load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN_person, checkpoint_person['model']) if args.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN_car, args.load_detectron) maskRCNN_person = mynn.DataParallel( maskRCNN_person, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN_person.eval() print('load cfg from file: {}'.format(args.cfg_file_car)) cfg_from_file(args.cfg_file_car) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) #cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights #assert_and_infer_cfg() maskRCNN_car = Generalized_RCNN() if args.cuda: maskRCNN_car.cuda() if args.load_ckpt_car: load_name = args.load_ckpt_car print("loading checkpoint for car %s" % (load_name)) checkpoint_car = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN_car, checkpoint_car['model']) maskRCNN_car = mynn.DataParallel(maskRCNN_car, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN_car.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) #imglist.sort(key=lambda f: int(filter(str.isdigit, f))) l = imglist try: imglist = sorted( l, key=lambda x: int(os.path.splitext(x)[0].split('/')[-1])) except: print('images couldnot be sorted') for i in xrange(num_images): print('img', i, ' out of ', num_images, ' filename: ', imglist[i].split('/')[-1], ' in camera', imglist[i].split('/')[-2]) im = cv2.imread(imglist[i]) try: assert im is not None except: continue timers = defaultdict(Timer) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) #print(im_name) output_name = os.path.basename(im_name) + '.txt' output_file = os.path.join(args.output_dir, '{}'.format(output_name)) text_file = open(output_file, "w") cfg_from_file(args.cfg_file_car) cls_boxes_car, cls_segms_car, cls_keyps_car, features_car = im_detect_all( maskRCNN_car, im, timers=timers) if len(cls_boxes_car[1]) > 0: features_car = features_car.data.cpu().numpy() #print(loop,loop2) #print(distance_matrix) #plt.figure() #plot_confusion_matrix(distance_matrix, classes=[0,1,2,3,4,5,6], #title='Confusion matrix, without normalization') #fig = plt.figure() #ax = fig.add_subplot(1,1,1) #ax.set_aspect('equal') #plt.imshow(distance_matrix, interpolation='nearest', cmap=plt.cm.ocean) #plt.colorbar() #plt.show() #fig.savefig('1.png') count = 0 filename = 'finalized_model.txt' loaded_model = pickle.load(open(filename, 'rb')) for ind, bb in enumerate(cls_boxes_car[1]): string = str(count) keyps = [k for klist in cls_keyps_car for k in klist] bb_new = [bb[0], bb[1], bb[2] - bb[0], bb[3] - bb[1]] features = features_car[ind, :, :, :].flatten() pca_feature = [] pca_feature.append(np.transpose(features.astype(np.float))) #print(pca_feature) features = loaded_model.transform(pca_feature) features = features[0] #loaded_model.transform(pca_feature) if bb[4] < 0.5: continue #for bb_ind,val in enumerate(bb_new): # string = string + ',' + str(val) for kp_ind, kp in enumerate(keyps[ind][0]): string = string + ',' + str(kp) + ',' + str( keyps[ind][1][kp_ind]) + ',' + str( int(keyps[ind][2][kp_ind])) for feature_ind, feature in enumerate(features): string = string + ',' + str(feature) string = string + ',car' text_file.write(string) text_file.write('\n') #print(string) count = count + 1 cfg_from_file(args.cfg_file_person) cls_boxes_person, cls_segms_person, cls_keyps_person, features_person = im_detect_all( maskRCNN_person, im, timers=timers) if len(cls_boxes_person[1]) > 0: features_person = features_person.data.cpu().numpy() for ind, bb in enumerate(cls_boxes_person[1]): string = str(count) keyps = [k for klist in cls_keyps_person for k in klist] bb_new = [bb[0], bb[1], bb[2] - bb[0], bb[3] - bb[1]] features = features_person[ind, :, :, :].flatten() pca_feature = [] pca_feature.append(np.transpose(features.astype(np.float))) #print(pca_feature) features = loaded_model.transform(pca_feature) features = features[0] #loaded_model.transform(pca_feature) #features = loaded_model.transform(np.transpose(features.astype(np.float))) # print(features) if bb[4] < 0.5: continue for bb_ind, val in enumerate(bb_new): string = string + ',' + str(val) for kp_ind, kp in enumerate(keyps[ind][0]): string = string + ',' + str(kp) + ',' + str( keyps[ind][1][kp_ind]) + ',' + str( int(keyps[ind][2][kp_ind])) for feature_ind, feature in enumerate(features): string = string + ',' + str(feature) string = string + ',person' text_file.write(string) text_file.write('\n') #print(string) count = count + 1 if save_image == True: im_name, _ = os.path.splitext(os.path.basename(imglist[i])) image_car = vis_utils.vis_one_image_car( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes_car, cls_segms_car, cls_keyps_car, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.5, kp_thresh=0.1) output_name = os.path.basename(im_name) + '.png' im = cv2.imread( os.path.join(args.output_dir, '{}'.format(output_name))) if im is None: continue continue vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes_person, cls_segms_person, cls_keyps_person, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.5, kp_thresh=10) if args.merge_pdfs and num_images > 1 and save_image == True: merge_out_path = '{}/results.pdf'.format(args.output_dir) if os.path.exists(merge_out_path): os.remove(merge_out_path) command = "pdfunite {}/*.pdf {}".format(args.output_dir, merge_out_path) subprocess.call(command, shell=True)