Exemplo n.º 1
0
 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)
Exemplo n.º 2
0
    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
Exemplo n.º 3
0
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")
Exemplo n.º 4
0
    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)
Exemplo n.º 6
0
    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()
Exemplo n.º 7
0
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
Exemplo n.º 8
0
    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()
Exemplo n.º 9
0
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
Exemplo n.º 10
0
    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))
Exemplo n.º 11
0
    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
Exemplo n.º 12
0
    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
Exemplo n.º 14
0
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
Exemplo n.º 15
0
    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
Exemplo n.º 16
0
    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
Exemplo n.º 17
0
 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)
Exemplo n.º 18
0
 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)
Exemplo n.º 19
0
 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"))
Exemplo n.º 20
0
    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()
Exemplo n.º 21
0
    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
Exemplo n.º 22
0
    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
Exemplo n.º 23
0
    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')
Exemplo n.º 24
0
    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
Exemplo n.º 26
0
    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
Exemplo n.º 27
0
 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)
Exemplo n.º 28
0
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
Exemplo n.º 29
0
  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
Exemplo n.º 30
0
    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)
Exemplo n.º 31
0
 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 {}
Exemplo n.º 32
0
 def _load_model(self, checkpoint):
     load_state_dict(self.model, checkpoint.pop("model"))