add resuming (#63)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/yolo/v8/classify/train.py

@@ -4,6 +4,7 @@ import torch
 from ultralytics.yolo import v8
 from ultralytics.yolo.data import build_classification_dataloader
 from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG, BaseTrainer
+from ultralytics.yolo.utils.modeling import get_model
 from ultralytics.yolo.utils.modeling.tasks import ClassificationModel
 
 
@@ -12,13 +13,13 @@ class ClassificationTrainer(BaseTrainer):
     def set_model_attributes(self):
         self.model.names = self.data["names"]
 
-    def load_model(self, model_cfg, weights, data):
+    def load_model(self, model_cfg, weights):
         # TODO: why treat clf models as unique. We should have clf yamls?
         if weights and not weights.__class__.__name__.startswith("yolo"):  # torchvision
             model = weights
         else:
-            model = ClassificationModel(model_cfg, weights, data["nc"])
-        ClassificationModel.reshape_outputs(model, data["nc"])
+            model = ClassificationModel(model_cfg, weights, self.data["nc"])
+        ClassificationModel.reshape_outputs(model, self.data["nc"])
         for m in model.modules():
             if not weights and hasattr(m, 'reset_parameters'):
                 m.reset_parameters()
@@ -28,6 +29,9 @@ class ClassificationTrainer(BaseTrainer):
             p.requires_grad = True  # for training
         return model
 
+    def load_ckpt(self, ckpt):
+        return get_model(ckpt)
+
     def get_dataloader(self, dataset_path, batch_size, rank=0, mode="train"):
         return build_classification_dataloader(path=dataset_path,
                                                imgsz=self.args.img_size,
@@ -46,6 +50,12 @@ class ClassificationTrainer(BaseTrainer):
         loss = torch.nn.functional.cross_entropy(preds, batch["cls"])
         return loss, loss
 
+    def check_resume(self):
+        pass
+
+    def resume_training(self, ckpt):
+        pass
+
 
 @hydra.main(version_base=None, config_path=DEFAULT_CONFIG.parent, config_name=DEFAULT_CONFIG.name)
 def train(cfg):

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)

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)

ultralytics/yolo/v8/segment/val.py

@@ -242,7 +242,7 @@ class SegmentationValidator(BaseValidator):
                               cls,
                               bboxes,
                               masks,
-                              paths,
+                              paths=paths,
                               fname=self.save_dir / f"val_batch{ni}_labels.jpg",
                               names=self.names)