Fixed yolo pretraine errors

InnopolisUni · Aug 27, 2024 · ccd28fe · ccd28fe
1 parent b0ab74a
commit ccd28fe
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 14 deletions.
diff --git a/innofw/core/integrations/ultralytics/ultralytics_adapter.py b/innofw/core/integrations/ultralytics/ultralytics_adapter.py
@@ -151,26 +151,27 @@ def train(self, data: UltralyticsDataModuleAdapter, ckpt_path=None):
             project="train",
             name=name,)
 
-        if ckpt_path is None:
-            self.opt.update(
-                device=self.device,
-                epochs=self.epochs,
-                imgsz=data.imgsz,
-                data=data.data,
-                workers=data.workers,
-                batch=data.batch_size,
-            )
-            self.model.train(**self.opt, **self.hyp)
-        else:
+        if ckpt_path is not None:
             try:
                 ckpt_path = TorchCheckpointHandler().convert_to_regular_ckpt(
-                    ckpt_path, inplace=False, dst_path=None
+                    ckpt_path, inplace=False, dst_path=None, set_epoch=0
                 )
                 self.opt.update(resume=str(ckpt_path))
-                self.model.train(**self.opt, **self.hyp)
+                self.model.ckpt["epoch"] = 0
+                self.model.ckpt_path = ckpt_path
             except Exception as e:
                 print(e)
 
+        self.opt.update(
+            device=self.device,
+            epochs=self.epochs,
+            imgsz=data.imgsz,
+            data=data.data,
+            workers=data.workers,
+            batch=data.batch_size,
+        )
+        self.model.train(**self.opt, **self.hyp)
+
         self.update_checkpoints_path()
 
     def predict(self, data: UltralyticsDataModuleAdapter, ckpt_path=None):

diff --git a/innofw/utils/checkpoint_utils/base_checkpoint_handler.py b/innofw/utils/checkpoint_utils/base_checkpoint_handler.py
@@ -85,9 +85,12 @@ def convert_to_regular_ckpt(
         ckpt_path: Path,
         dst_path: Optional[Path] = None,
         inplace: bool = True,
+        set_epoch: int = -1
     ) -> Path:
         model = self.load_model(None, ckpt_path)
-
+        if set_epoch != -1:
+            if "epoch" in model.keys():
+                model["epoch"] = set_epoch
         if inplace:
             tmp_path = Path(tempfile.mkdtemp()) / ckpt_path.name
             self.save_ckpt(model, tmp_path, None, wrap=False)