如何使用 Detectron2 .pth 模型从存储中进行预测..？

Posted 2023-03-27

【中文标题】如何使用 Detectron2 .pth 模型从存储中进行预测..？

【英文标题】：How to use the Detectron2 .pth model for prediction from storage..?

【发布时间】：2020-12-02 22:22:01

【问题描述】：

我已经在 google colab server 免费服务器上训练了 Detectron2 模型。

from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
import os

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO- 
InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("Dataset_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO- 
InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 2800
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

上面的代码创建了一个“输出”文件夹，其中有 4 个文件： model_final.pth、metrics.json、last_checkpoint 和 events.out.file

我可以使用这个模型进行预测

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.DATASETS.TEST = ("Datset_test")
predictor = DefaultPredictor(cfg)

colab 本地会话未过期时一切正常。

问题：

当我将此输出文件夹安装在未导入模型之外的其他位置时。 如何使用 .pth 保存的模型进行预测？ 如何重用经过detectron2 训练的模型进行预测..？

【参考方案1】：

下载您的 .pth 文件。当您打开一个新的 Colab 运行时，将该文件上传到您的本地运行时（您可以使用“！cp”从 Google Drive 到您的运行时快速传输 .pth 文件）然后按照官方教程中的步骤进行操作：https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5

你需要使用你的配置文件cfg.merge_from_file("YOUR CONFIG FILE")和你的.pth文件路径：cfg.MODEL.WEIGHTS = ("YOUR .PTH FILE, PROBABLY /CONTENT/MODEL_FINAL.PTH")