0

I have trained the Detectron2 model on the google colab server free server.

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

This above code creates an "output" folder in which I have 4 files: model_final.pth, metrics.json, last_checkpoint and events.out.file

model_final.pth, metrics.json, last_checkpoint and events.out.file

I can use this model for prediction using

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)

Everything works fine when colab local session is not expired.

Problem:

  1. When I mount this output folder somewhere else than the model not imported.
  2. How to use .pth saved model for predictions?
  3. How to reuse detectron2 trained model for prediction..?
JB.py
  • 59
  • 2
  • 8

1 Answers1

0

Download your .pth file. When you open a new Colab runtime, upload that file to your local runtime (you can use "!cp" from Google Drive to your runtime to transfer .pth file quickly) and then follow the steps on official tutorial : https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5

You need to use your config file cfg.merge_from_file("YOUR CONFIG FILE") and your .pth file path: cfg.MODEL.WEIGHTS = ("YOUR .PTH FILE, PROBABLY /CONTENT/MODEL_FINAL.PTH")

knayman
  • 97
  • 8