如何在 Detectron2 中保存和加载自定义数据集的模型？答案

【问题标题】：How to save and load models of custom dataset in Detectron2?如何在 Detectron2 中保存和加载自定义数据集的模型？
【发布时间】：2020-07-30 04:21:55
【问题描述】：

我尝试使用以下方法保存和加载模型：所有键都已映射，但输出中没有预测 #1

from detectron2.modeling import build_model
model = build_model(cfg)
torch.save(model.state_dict(), 'checkpoint.pth') 
model.load_state_dict(torch.load(checkpoint_path,map_location='cpu'))

我也尝试过使用官方文档但无法理解输入格式部分

from detectron2.checkpoint import DetectionCheckpointer
DetectionCheckpointer(model).load(file_path_or_url)  # load a file, usually from cfg.MODEL.WEIGHTS
checkpointer = DetectionCheckpointer(model, save_dir="output")
checkpointer.save("model_999")  # save to output/model_999.pth

【问题讨论】：

你能分享错误和更多代码让我们重现错误吗？

标签： deep-learning pytorch detectron

【解决方案1】：

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file('COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml'))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # Set threshold for this model
cfg.MODEL.WEIGHTS = '/content/model_final.pth' # Set path model .pth
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
predictor = DefaultPredictor(cfg)

我加载自定义模型的代码有效。

【讨论】：