add resuming (#63)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2022-12-05 20:56:41 -06:00
parent de3e6ca54d
commit fbeeb5d1e1
7 changed files with 86 additions and 30 deletions
--- a/ultralytics/yolo/v8/detect/train.py
+++ b/ultralytics/yolo/v8/detect/train.py
@ -15,10 +15,10 @@ from .val import DetectionValidator
 # BaseTrainer python usage
 class DetectionTrainer(SegmentationTrainer):

-    def load_model(self, model_cfg, weights, data):
+    def load_model(self, model_cfg, weights):
        model = DetectionModel(model_cfg or weights["model"].yaml,
                               ch=3,
-                               nc=data["nc"],
+                               nc=self.data["nc"],
                               anchors=self.args.get("anchors"))
        if weights:
            model.load(weights)