haungjj:load pretrained detector for mix up model (facebookresearch#2)

2333

Author: hjjpku

maskrcnn_benchmark/huangjj.py

@@ -0,0 +1,46 @@
+# huangjj-pkusz-pcl
+# load the pretrain fpn(it should be a checkpoint, and the last checkpoint is loaded as default) to initialize the mix-up model
+import logging
+import torch
+import os
+from maskrcnn_benchark.utils.model_serialization import load_state_dict
+
+
+def get_checkpoint_file(path):
+	save_file = path
+	try:
+	    with open(save_file, "r") as f:
+	        last_saved = f.read()
+	        last_saved = last_saved.strip()
+	except IOError:
+	    # if file doesn't exist, maybe because it has just been
+	    # deleted by a separate process
+	    last_saved = ""
+	return last_saved
+
+def load_file(cfg,f):
+	if f.endswith(".pkl"):
+		return load_c2_format(cfg, f)
+	return torch.load(f, map_location=torch.device("cpu"))
+
+def load_model(model, checkpoint)
+	load_state_dict(model, checkpoint.pop("model"))
+	return model
+
+def load_pretrain_detector(cfg, model)
+	logger = logging.getLogger(__name__)
+
+	path_to_model = cfg.PRETRAIN_DET_DIR
+	save_file = os.path.join(path_to_model, "last_checkpoint")
+
+	if os.path.exists(save_file):
+		f = get_model_file(path_to_model) 	
+	if not f:
+		logger.infor("No pretrained detecotr found.")
+		exit(0)
+	
+	checkpoint = load_file(cfg,f)
+	
+	return load_model(model, checkpoint)
+	
+