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"))