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 · 2022-12-05 20:56:41 -06:00 · fbeeb5d1e1
commit fbeeb5d1e1
parent de3e6ca54d
7 changed files with 86 additions and 30 deletions
--- a/ultralytics/yolo/v8/segment/train.py
+++ b/ultralytics/yolo/v8/segment/train.py
@ -26,10 +26,10 @@ class SegmentationTrainer(BaseTrainer):
        batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
        return batch

-    def load_model(self, model_cfg, weights, data):
+    def load_model(self, model_cfg, weights):
        model = SegmentationModel(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)