def _load_model(self, checkpoint): ckpt_model = checkpoint.pop("model") if isinstance(ckpt_model, torch.Tensor): # self.model = ckpt_model pass else: load_state_dict(self.model, ckpt_model)
def _build_detection_model(self): """ Build the detection model. Builds a CNN using the vqa-maskrcnn-benchmark repository. """ from maskrcnn_benchmark.config import cfg from maskrcnn_benchmark.modeling.detector import build_detection_model from maskrcnn_benchmark.utils.model_serialization import load_state_dict dp = self.opt['datapath'] build(dp) cfg_path = os.path.join(dp, 'models/detectron/detectron_config.yaml') model_path = os.path.join(dp, 'models/detectron/detectron_model.pth') cfg.merge_from_file(cfg_path) cfg.freeze() model = build_detection_model(cfg) checkpoint = torch.load(model_path, map_location=torch.device("cpu")) load_state_dict(model, checkpoint.pop("model")) if self.use_cuda: model.to("cuda") model.eval() self.detection_model = model
def main(): parser = init() args = parser.parse_args() # update the config options with the config file cfg.merge_from_file(args.config_file) # manual override some options cfg.merge_from_list(["MODEL.DEVICE", "cpu"]) # cfg.merge_from_list(["MODEL.MASK_ON", False]) coco_demo = COCODemo( cfg, confidence_threshold=args.confidence_threshold, show_mask_heatmaps=args.show_mask_heatmaps, masks_per_dim=args.masks_per_dim, min_image_size=args.min_image_size, ) if args.weight is not None: checkpoint = torch.load(args.weight) load_state_dict(coco_demo.model, checkpoint.pop("model")) del checkpoint start_time = time.time() img = cv2.imread(args.img) composite = coco_demo.run_on_opencv_image(img) print("Time: {:.2f} s / img".format(time.time() - start_time)) # plt.figure() plt.imshow(composite) plt.savefig("result.jpg")
def __init__(self, exp_dict): super().__init__() cfg_base_path = "./models/configs/" self.n_classes = 21 cfg_path = cfg_base_path + "e2e_mask_rcnn_R_50_FPN_1x.yaml" self.cfg = cfg self.cfg.merge_from_file(cfg_path) self.cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES = self.n_classes # --------------- # build model self.backbone_fpn = backbone.build_backbone(self.cfg) self.rpn = rpn.build_rpn(self.cfg, self.backbone_fpn.out_channels) self.roi_heads = roi_heads.build_roi_heads( self.cfg, self.backbone_fpn.out_channels) # --------------- # load checkpoint checkpoint = _load_file(self.cfg) load_state_dict(self, checkpoint.pop("model")) #-------- # Opt stage self.cfg.SOLVER.BASE_LR = ((0.0025 * 8) / (16 / float(exp_dict["batch_size"]))) optimizer = make_optimizer(self.cfg, self) scheduler = make_lr_scheduler(self.cfg, optimizer) self.opt = optimizer self.scheduler = scheduler
def __init__(self, config_file, weight_file, min_image_size=(800, ), dataset='coco', spire_dir='/tmp', class_transfor=None, origin_size=False, logger=logging.getLogger()): # 加载模型配置文件 cfg.merge_from_file(config_file) cfg.freeze() self.cfg = cfg model = build_detection_model(cfg) self.device = torch.device(cfg.MODEL.DEVICE) model.to(self.device) # 读取与加载网络参数 checkpoint = torch.load(weight_file, map_location=torch.device("cpu")) load_state_dict(model, checkpoint.pop("model")) model.eval() self.model = model if origin_size: self.transforms = [build_transform_origin(cfg)] else: self.transforms = [ build_transform(cfg, _size) for _size in min_image_size ] self.spire_anno = SpireAnno(dataset=dataset, spire_dir=spire_dir)
def __init__(self): super(VQAMaskRCNNBenchmark, self).__init__() self.model = build_detection_model(cfg) model_state_dict = torch.load(maskrcnn_checkpoint) load_state_dict(self.model, model_state_dict.pop("model")) # make sure maskrcnn_benchmark is in eval mode self.model.eval()
def train(cfg): model = build_detection_model(cfg) print(model) device = torch.device(cfg.MODEL.DEVICE) model.to(device) model2 = rating_model(cfg) model2.to(cfg.MODEL.DEVICE) print(model2) optimizer = make_optimizer(cfg, model, model2) scheduler = make_lr_scheduler(cfg, optimizer) save_to_disk = get_rank() == 0 checkpointer1 = DetectronCheckpointer(cfg, model, optimizer, scheduler, cfg.OUTPUT_DIR, save_to_disk) checkpointer2 = DetectronCheckpointer(cfg, model2, optimizer, scheduler, cfg.OUTPUT_DIR, save_to_disk) backbone_parameters = torch.load(cfg.MODEL.WEIGHT, map_location=torch.device("cpu")) newdict = {} newdict['model'] = removekey(backbone_parameters['model'], [ 'module.roi_heads.box.predictor.cls_score.bias', 'module.roi_heads.box.predictor.cls_score.weight', 'module.roi_heads.box.predictor.bbox_pred.bias', 'module.roi_heads.box.predictor.bbox_pred.weight' ]) load_state_dict(model, newdict.pop("model")) arguments = {} arguments["iteration"] = 0 checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD data_loader = make_data_loader( cfg, is_train=True, is_distributed=False, start_iter=arguments["iteration"], ) proposing_train( model, model2, data_loader, optimizer, scheduler, checkpointer1, checkpointer2, device, checkpoint_period, arguments, ) return model, model2
def __init__(self): super(VQAMaskRCNNBenchmark, self).__init__() # self.avgpool = nn.AdaptiveAvgPool2d(output_size=1) self.model = build_detection_model(cfg) model_state_dict = torch.load(maskrcnn_checkpoint) load_state_dict(self.model, model_state_dict.pop("model")) # make sure maskrcnn_benchmark is in eval mode self.model.eval()
def dla(cfg): model = BACKBONE[cfg.MODEL.BACKBONE.CONV_BODY]( cfg.MODEL.DLA.STAGE_WITH_DCN) # Load the ImageNet pretrained backbone if no valid pre-trained model weights are given if not os.path.exists(cfg.MODEL.WEIGHT): state_dict = load_state_dict_from_url( model_urls[BACKBONE_ARCH[cfg.MODEL.BACKBONE.CONV_BODY]], progress=True) load_state_dict(model, state_dict) return model
def test_complex_model_loaded(self): for add_data_parallel in [False, True]: model, state_dict = self.create_complex_model() if add_data_parallel: model = nn.DataParallel(model) load_state_dict(model, state_dict) for loaded, stored in zip(model.state_dict().values(), state_dict.values()): # different tensor references self.assertFalse(id(loaded) == id(stored)) # same content self.assertTrue(loaded.equal(stored))
def _build_detection_model(self): cfg.merge_from_file(self.args.config_file) cfg.freeze() model = build_detection_model(cfg) checkpoint = torch.load(self.args.model_file, map_location=torch.device("cpu")) load_state_dict(model, checkpoint.pop("model")) model.to("cuda") model.eval() return model
def test_complex_model_loaded(self): for add_data_parallel in [False, True]: model, state_dict = self.create_complex_model() if add_data_parallel: model = nn.DataParallel(model) load_state_dict(model, state_dict) for loaded, stored in zip(model.state_dict().values(), state_dict.values()): # different tensor references self.assertFalse(id(loaded) == id(stored)) # same content self.assertTrue(loaded.equal(stored))
def load_detection_model(yaml_file, yaml_ckpt): cfg.merge_from_file(yaml_file) cfg.freeze() model = build_detection_model(cfg) checkpoint = torch.load(yaml_ckpt, map_location=torch.device("cpu")) load_state_dict(model, checkpoint.pop("model")) model.to("cuda") model.eval() return model
def build_mask_rcnn(): ckpt_dir = Path('./mask_rcnn_ckpt').resolve() config_path = ckpt_dir.joinpath('detectron_config.yaml') model_path = ckpt_dir.joinpath('detectron_model.pth') cfg.merge_from_file(str(config_path)) cfg.freeze() model = build_detection_model(cfg) checkpoint = torch.load(str(model_path), map_location=torch.device("cpu")) load_state_dict(model, checkpoint.pop("model")) return model
def _build_detection_model(self): cfg.merge_from_file('model_data/detectron_model.yaml') cfg.freeze() model = build_detection_model(cfg) checkpoint = torch.load('model_data/detectron_model.pth', map_location=torch.device("cpu")) load_state_dict(model, checkpoint.pop("model")) model.to("cuda") model.eval() return model
def _build_detection_model(self, checkpoint_path, config_path): cfg.merge_from_file(config_path) cfg.freeze() model = build_detection_model(cfg) checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu")) load_state_dict(model, checkpoint.pop("model")) model.to(self.device) model.eval() return model
def _load_model(self, checkpoint): model_dict = self.model.state_dict() ''' for k in model_dict.keys(): print(k) print(checkpoint.keys()) for k in checkpoint['model']['net'].keys(): print(k) ''' checkpoint = checkpoint.pop("model") if 'net' in checkpoint: load_state_dict(self.model, checkpoint.pop("net")) else: load_state_dict(self.model, checkpoint)
def _load_model_except_box_branch(self, checkpoint): if "model" in checkpoint: saved_state_dict = checkpoint.pop('model') else: saved_state_dict = checkpoint new_params = self.model.state_dict().copy() for i in saved_state_dict: # 'module.roi_heads.box.xx' i_parts = i.split('.') if i_parts[2] == 'box': continue new_params['.'.join(i_parts)] = saved_state_dict[i] print('.'.join(i_parts)) load_state_dict(self.model, new_params)
def _load_model(self, checkpoint): # ################ add by hui ######################## models = {} for key, value in checkpoint['model'].items(): is_ignore_key = False for p_key in self.ignore_keys: if p_key in key: is_ignore_key = True if not is_ignore_key: models[key] = value else: self.logger.info("ignore {}".format(key)) checkpoint['model'] = models # ################################################### load_state_dict(self.model, checkpoint.pop("model"))
def __init__(self, cfg_path, weights_path, input_shape=(608, 608)): cfg.merge_from_file(cfg_path) cfg.merge_from_list(['DTYPE', 'float16']) self._cfg = cfg.clone() self._model = build_detection_model(self._cfg) self._model.eval() self._device = 'cuda' self._model.to(self._device) self.shape = input_shape save_dir = cfg.OUTPUT_DIR checkpoint = torch.load(weights_path, map_location=torch.device("cpu")) load_state_dict(self._model, checkpoint.pop("model")) self._transform = self._build_transform()
def _build_detection_model(self): cfg.merge_from_file( '/home/leonardo/hack_fb/content/model_data/detectron_model.yaml') cfg.freeze() model = build_detection_model(cfg) checkpoint = torch.load( '/home/leonardo/hack_fb/content/model_data/detectron_model.pth', map_location='cpu') load_state_dict(model, checkpoint.pop("model")) model.to(torch.device("cpu")) model.eval() return model
def load(self, f=None, test=False): if test: self.logger.info("Loading checkpoint from {}".format(f)) checkpoint = self._load_file(f) load_state_dict(self.model, checkpoint.pop("model")) else: if 'e2e_mask_rcnn_R_50_FPN_1x.pth' in f: self.transfer_learning = True else: self.transfer_learning = False if not self.transfer_learning and self.has_checkpoint(): # override argument with existing checkpoint self.transfer_learning = False f = self.get_checkpoint_file() # f = False if not f: # no checkpoint could be found self.logger.info( "No checkpoint found. Initializing model from scratch") return {} self.logger.info("Loading checkpoint from {}".format(f)) checkpoint = self._load_file(f) # delete this two because we add new module which transfer learning model does not have del checkpoint['scheduler'], checkpoint['optimizer'] self._load_model(checkpoint) if self.transfer_learning: # default last epoch of loaded weight is 89999 checkpoint['iteration'] = -1 if not self.transfer_learning: # if use transfer learning , do not load pretrain model scheduler and optimizer if "optimizer" in checkpoint and self.optimizer: self.logger.info("Loading optimizer from {}".format(f)) # pdb.set_trace() # pdb.set_trace() self.optimizer.load_state_dict(checkpoint.pop("optimizer")) if "scheduler" in checkpoint and self.scheduler: self.logger.info("Loading scheduler from {}".format(f)) self.scheduler.load_state_dict(checkpoint.pop("scheduler")) # return any further checkpoint data return checkpoint
def _load_model(self, checkpoint): loaded = checkpoint.pop('model') unload_keyword = list(self.cfg.FEW_SHOT.UNLOAD_KEYWORD) if self.cfg.FEW_SHOT.LOAD_PRETRIANED_RPN_ONLY: unload_keyword += ['roi_heads'] for key, value in self.model.state_dict().items(): if any([keyword in key for keyword in unload_keyword]): if 'module' not in key: key = 'module.' + key loaded[key] = value load_state_dict(self.model, loaded) self.logger.info( 'loading module without keyword {}'.format(unload_keyword)) self.logger.info('successfully loading model')
def build_detect_model(self): print('Tiki : Initializing : Building - Detectron') cfg.merge_from_file('/final/data/detectron_model.yaml') cfg.freeze() model = build_detection_model(cfg) chkpt = torch.load('/final/data/detectron_model.pth', map_location=self.device) load_state_dict(model, chkpt.pop('model')) model.to(self.device.type) model.eval() return model
def _build_detection_model(self): from maskrcnn_benchmark.utils.model_serialization import load_state_dict from maskrcnn_benchmark.modeling.detector import build_detection_model from maskrcnn_benchmark.config import cfg cfg.merge_from_file(self.cfg_file) cfg.freeze() model = build_detection_model(cfg) checkpoint = torch.load(self.model_file, map_location=self.device) load_state_dict(model, checkpoint.pop("model")) model.to(self.device) model.eval() print(self.__class__.__name__, ": Loaded Model...") return model
def _build_detection_model(self): cfg.merge_from_file( os.path.join(BASE_VQA_DIR_PATH, "model_data/detectron_model.yaml")) cfg.freeze() model = build_detection_model(cfg) checkpoint = torch.load( os.path.join(BASE_VQA_DIR_PATH, "model_data/detectron_model.pth"), map_location=torch.device("cpu"), ) load_state_dict(model, checkpoint.pop("model")) model.to("cuda") model.eval() return model
def _load_model(self, checkpoint): if not self.transfer_learning: load_state_dict(self.model, checkpoint.pop("model")) else: # pdb.set_trace() # delete roi_head.box/mask.predictor.cls_score/bbox_pred/mask_fcn_logits in state_dict pretrained_weights = checkpoint.pop("model") model_state_dict = self.model.state_dict() loaded_state_dict = strip_prefix_if_present(pretrained_weights, prefix="module.") align_and_update_state_dicts(model_state_dict, loaded_state_dict) model_state_dict = { k: v for k, v in model_state_dict.items() if 'cls_score' not in k and 'bbox_pred' not in k and 'mask_fcn_logits' not in k } self.model.load_state_dict(model_state_dict, strict=False)
def train(cfg): model = build_detection_model(cfg) device = torch.device(cfg.MODEL.DEVICE) model.to(device) print(model) model3 = predicate_model(cfg) model3.to(cfg.MODEL.DEVICE) print(model3) optimizer = make_optimizer(cfg, model3) scheduler = make_lr_scheduler(cfg, optimizer) save_to_disk = get_rank() == 0 checkpointer = DetectronCheckpointer(cfg, model3, optimizer, scheduler, cfg.OUTPUT_DIR, save_to_disk) backbone_parameters = torch.load(os.getcwd() + cfg.CONFIG.backbone_weight, map_location=torch.device("cpu")) newdict = {} newdict['model'] = removekey(backbone_parameters['model'], []) load_state_dict(model, newdict.pop("model")) arguments = {} arguments["iteration"] = 0 checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD data_loader = make_data_loader( cfg, is_train=True, is_distributed=False, start_iter=arguments["iteration"], ) predicate_train( model, model3, data_loader, optimizer, scheduler, checkpointer, device, checkpoint_period, arguments, ) return model, model3
def _build_detection_model(self): # detectron_model.yaml의 주소 # cfg.merge_from_file('/home/multicam/multicam_project/speak_image/IC/model_data/detectron_model.yaml') cfg.merge_from_file('/home/ubuntu/test/backend/Codes/speak_image/IC/model_data/detectron_model.yaml') cfg.freeze() model = build_detection_model(cfg) # detectron_model.pth의 주소 # checkpoint = torch.load('C:\\Users\\multicampus\\subproject3\\backend\\Codes\\checkpoints\\detectron_model.pth', # map_location=torch.device("cuda")) checkpoint = torch.load('/home/ubuntu/test/IC/detectron_model.pth', map_location=torch.device("cpu")) load_state_dict(model, checkpoint.pop("model")) model.to("cpu") model.eval() return model
def _load_model_except_class_layer(self, checkpoint): if "model" in checkpoint: saved_state_dict = checkpoint['model'] else: saved_state_dict = checkpoint new_params = self.model.state_dict().copy() for i in saved_state_dict: # 'module.roi_heads.mask.predictor.mask_fcn_logits.bias' # 'module.roi_heads.box.predictor.cls_score.weight' # 'module.roi_heads.box.predictor.bbox_pred.weight' i_parts = i.split('.') if (i_parts[2] == 'mask' and i_parts[4] == 'mask_fcn_logits') or \ (i_parts[2] == 'box' and i_parts[4] == 'cls_score') or \ (i_parts[2] == 'box' and i_parts[4] == 'bbox_pred'): continue new_params['.'.join(i_parts)] = saved_state_dict[i] print('.'.join(i_parts)) load_state_dict(self.model, new_params)
def load_merge_checkpoints(self, app_file, sp_file): self.logger.info("Loading app checkpoint from {}".format(app_file)) self.logger.info("Loading sp checkpoint from {}".format(sp_file)) checkpoint_app = torch.load(app_file, map_location=torch.device("cpu")) checkpoint_sp = torch.load(sp_file, map_location=torch.device("cpu")) weights_app = checkpoint_app['model'] new_dict_app = { k.replace('module.', ''): v for k, v in weights_app.items() if 'sp_branch' not in k } weights_sp = checkpoint_sp['model'] new_dict_sp = { k.replace('module.', ''): v for k, v in weights_sp.items() if 'sp_branch' in k } this_state = self.model.state_dict() this_state.update(new_dict_app) this_state.update(new_dict_sp) # self.model.load_state_dict(this_state) load_state_dict(self.model, this_state) # return any further checkpoint data return {}
def _load_model(self, checkpoint): load_state_dict(self.model, checkpoint.pop("model"))