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