opensportslib 0.1.2.dev6__tar.gz → 0.1.2.dev7__tar.gz

{opensportslib-0.1.2.dev6/opensportslib.egg-info → opensportslib-0.1.2.dev7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opensportslib
-Version: 0.1.2.dev6
+Version: 0.1.2.dev7
 Summary: OpenSportsLib is the professional library, designed for advanced video understanding in sports. It provides state-of-the-art tools for action recognition, spotting, retrieval, and captioning, making it ideal for researchers, analysts, and developers working with sports video data.
 Author: Jeet Vora
 Requires-Python: >=3.12

{opensportslib-0.1.2.dev6 → opensportslib-0.1.2.dev7}/opensportslib/apis/localization.py

@@ -15,6 +15,8 @@ class LocalizationModel(BaseTaskModel):
             self.last_loaded_weights = weights
             self.best_checkpoint = weights
 
+        self.train_flag = False  # Flag to indicate whether we're in training mode (affects checkpoint loading behavior)
+
     def _resolve_split_path(self, split: str, override: str | None = None) -> str:
         if override is not None:
             return expand(override)
@@ -69,15 +71,18 @@ class LocalizationModel(BaseTaskModel):
             load_checkpoint,
             localization_remap,
         )
-
+        from opensportslib.core.optimizer.builder import build_optimizer
+        from opensportslib.core.scheduler.builder import build_scheduler
+        default_args = kwargs.get("default_args", None)
         del kwargs
         if weights is None:
             raise ValueError("`weights` must be provided to load_weights().")
 
         model_cfg = getattr(self.config, "MODEL", None)
-        original_multi_gpu = getattr(model_cfg, "multi_gpu", None)
-        if model_cfg is not None and original_multi_gpu is not None:
-            model_cfg.multi_gpu = False
+        if not self.train_flag:
+            original_multi_gpu = getattr(model_cfg, "multi_gpu", None)
+            if model_cfg is not None and original_multi_gpu is not None:
+                model_cfg.multi_gpu = False
 
         device = select_device(self.config.SYSTEM)
         if self.model is None:
@@ -90,9 +95,28 @@ class LocalizationModel(BaseTaskModel):
         if is_local_path(weights):
             self.config.SYSTEM.work_dir = os.path.dirname(os.path.abspath(weights))
 
-        inner_model, _, _, _ = load_checkpoint(
+        if default_args is not None:
+            logging.info("Building optimizer + scaler for checkpoint restore...")
+            optimizer, scaler = build_optimizer(
+                inner_model.parameters(),  # or _get_params() if required
+                self.config.TRAIN.optimizer
+            )
+            
+            logging.info("Building scheduler for checkpoint restore...")
+            scheduler = build_scheduler(
+                optimizer,
+                self.config.TRAIN.scheduler,
+                default_args
+            )
+        else:
+            optimizer = scheduler = scaler = None
+
+        inner_model, optimizer, scheduler, scaler, epoch, checkpoint = load_checkpoint(
             model=inner_model,
             path=weights,
+            optimizer=optimizer,
+            scheduler=scheduler,
+            scaler=scaler,
             device=device,
             key_remap_fn=localization_remap,
         )
@@ -107,8 +131,24 @@ class LocalizationModel(BaseTaskModel):
         self.last_loaded_weights = weights
         self.best_checkpoint = weights
 
-        if model_cfg is not None and original_multi_gpu is not None:
-            model_cfg.multi_gpu = original_multi_gpu
+        best_epoch = checkpoint.get("best_epoch", 0)
+
+        best_criterion_valid = checkpoint.get(
+            "best_criterion_valid",
+            0 if self.config.TRAIN.criterion_valid == "map" else float("inf")
+        )
+        self._resume_state = {
+            "optimizer": optimizer,
+            "scheduler": scheduler,
+            "scaler": scaler,
+            "epoch": epoch if epoch is not None else 0,
+            "best_epoch": best_epoch,
+            "best_criterion_valid": best_criterion_valid,
+        }
+
+        if not self.train_flag:
+            if model_cfg is not None and original_multi_gpu is not None:
+                model_cfg.multi_gpu = original_multi_gpu
 
     def train(
         self,
@@ -167,13 +207,6 @@ class LocalizationModel(BaseTaskModel):
 
         start = time.time()
 
-        if effective_weights is not None:
-            if self.model is None or self.last_loaded_weights != effective_weights:
-                self.load_weights(weights=effective_weights)
-        elif self.model is None:
-            device = select_device(self.config.SYSTEM)
-            self.model = build_model(self.config, device=device)
-
         data_obj_train = build_dataset(self.config, split="train")
         dataset_train = data_obj_train.building_dataset(
             cfg=data_obj_train.cfg,
@@ -200,11 +233,21 @@ class LocalizationModel(BaseTaskModel):
             dali=self.config.dali,
         )
 
+        default_args = get_default_args_trainer(self.config, len(train_loader))
+        
+        self.train_flag = True  # Set flag to indicate training mode for checkpoint loading
+        if effective_weights is not None:
+            if self.model is None or self.last_loaded_weights != effective_weights:
+                self.load_weights(weights=effective_weights, default_args=default_args)
+        elif self.model is None:
+            device = select_device(self.config.SYSTEM)
+            self.model = build_model(self.config, device=device)
+
         self.trainer = build_trainer(
             cfg=self.config,
             model=self.model,
-            default_args=get_default_args_trainer(self.config, len(train_loader)),
-            resume_from=effective_weights,
+            default_args=default_args,
+            resume_from=self._resume_state if hasattr(self, "_resume_state") else None,
         )
 
         logging.info("Start training")

{opensportslib-0.1.2.dev6 → opensportslib-0.1.2.dev7}/opensportslib/core/trainer/classification_trainer.py

@@ -1167,11 +1167,12 @@ class Trainer_Classification:
             from opensportslib.models.builder import build_model
             if self.model is None:
                 self.model, _ = build_model(self.config, self.device)
-            self.model, optimizer, scheduler, epoch = load_checkpoint(
+            self.model, optimizer, scheduler, scaler, epoch, checkpoint = load_checkpoint(
                 self.model, path, optimizer, scheduler, device=self.device
             )
             self.optimizer = optimizer
             self.scheduler = scheduler
+            self.scaler = scaler
             self.epoch = epoch
             logging.info(f"Model loaded from {path}, epoch: {epoch}")
             return self.model, self.optimizer, self.scheduler, self.epoch

{opensportslib-0.1.2.dev6 → opensportslib-0.1.2.dev7}/opensportslib/core/trainer/localization_trainer.py

@@ -29,7 +29,6 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 """
 from opensportslib.metrics.localization_metric import *
 from opensportslib.core.optimizer.builder import build_optimizer
-from opensportslib.core.optimizer.builder import build_optimizer
 from opensportslib.core.scheduler.builder import build_scheduler
 from opensportslib.core.utils.config import store_json
 from opensportslib.datasets.builder import build_dataset
@@ -67,20 +66,10 @@ def build_trainer(cfg, model=None, default_args=None, resume_from=None):
         
         # Handle checkpoint loading
         if resume_from is not None:
-            if not os.path.isfile(resume_from):
-                raise ValueError(f"Checkpoint file not found: {resume_from}")
-                
-            logging.info(f"Loading checkpoint from: {resume_from}")
-            checkpoint = torch.load(resume_from)
-            
-            # Load model state
-            model.load(checkpoint['model_state_dict'])
-            logging.info("Model state loaded successfully")
-            
-            # Get current training progress
-            start_epoch = checkpoint['epoch'] + 1
-            logging.info(f"Resuming from epoch {start_epoch}")
-            
+            optimizer = resume_from["optimizer"]
+            scheduler = resume_from["scheduler"]
+            scaler = resume_from["scaler"]
+            start_epoch = resume_from["epoch"] + 1
             # Check if we've already reached target epochs
             if start_epoch >= cfg.TRAIN.num_epochs:
                 logging.error(f"Model already trained for {start_epoch} epochs")
@@ -89,38 +78,18 @@ def build_trainer(cfg, model=None, default_args=None, resume_from=None):
                 raise ValueError("Need to increase num_epochs to continue training")
             
             logging.info(f"Will continue training from epoch {start_epoch} to {cfg.TRAIN.num_epochs}")
-        
-        logging.info("Building optimizer...")
-        optimizer, scaler = build_optimizer(model._get_params(), cfg.TRAIN.optimizer)
-        
-        # Load optimizer state if available in checkpoint
-        if resume_from is not None and 'optimizer_state_dict' in checkpoint:
-            try:
-                optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-                scaler.load_state_dict(checkpoint['scaler_state_dict'])
-                logging.info("Optimizer and scaler states loaded")
-            except Exception as e:
-                logging.warning(f"Could not load optimizer state: {e}")
-                logging.warning("Will start with fresh optimizer state")
-            
-        logging.info("Building scheduler...")
-        lr_scheduler = build_scheduler(optimizer, cfg.TRAIN.scheduler, default_args)
-        
-        # Load scheduler state if available
-        if resume_from is not None and 'lr_state_dict' in checkpoint:
-            try:
-                lr_scheduler.load_state_dict(checkpoint['lr_state_dict'])
-                logging.info("Scheduler state loaded")
-            except Exception as e:
-                logging.warning(f"Could not load scheduler state: {e}")
-                logging.warning("Will start with fresh scheduler state")
+        else:
+            logging.info("Building optimizer...")
+            optimizer, scaler = build_optimizer(model._get_params(), cfg.TRAIN.optimizer)
+            logging.info("Building scheduler...")
+            scheduler = build_scheduler(optimizer, cfg.TRAIN.scheduler, default_args)
 
         trainer = Trainer_e2e(
             cfg,
             model,
             optimizer,
             scaler,
-            lr_scheduler,
+            scheduler,
             default_args["work_dir"],
             default_args["dali"],
             default_args["repartitions"],
@@ -132,8 +101,8 @@ def build_trainer(cfg, model=None, default_args=None, resume_from=None):
         
         # Load training history if resuming
         if resume_from is not None:
-            trainer.best_epoch = checkpoint.get('best_epoch', 0)
-            trainer.best_criterion_valid = checkpoint.get('best_criterion_valid', 
+            trainer.best_epoch = resume_from.get('best_epoch', 0)
+            trainer.best_criterion_valid = resume_from.get('best_criterion_valid', 
                 0 if cfg.TRAIN.criterion_valid == "map" else float("inf"))
             logging.info(f"Restored best epoch: {trainer.best_epoch}")
     
@@ -441,7 +410,7 @@ class Trainer_e2e(Trainer):
         best_checkpoint_path = os.path.join(
             self.save_dir, f"best_checkpoint.pt"
         )
-        self.model._model, _, _, epoch = load_checkpoint(model=self.model._model,
+        self.model._model, _, _, _, epoch, _ = load_checkpoint(model=self.model._model,
                                         path=best_checkpoint_path,
                                         key_remap_fn=localization_remap)
         logging.info(f"Loaded best model from epoch {self.best_epoch}")

{opensportslib-0.1.2.dev6 → opensportslib-0.1.2.dev7}/opensportslib/core/utils/checkpoint.py

@@ -76,6 +76,7 @@ def load_checkpoint(
     path,
     optimizer=None,
     scheduler=None,
+    scaler=None,
     device=None,
     key_remap_fn=None,
     hf_filename="model.pth.tar",   # required if loading from HF repo
@@ -164,7 +165,7 @@ def load_checkpoint(
     # --------------------------------------------------
     # Load checkpoint
     # --------------------------------------------------
-    checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
+    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
 
     # ---------------- MODEL STATE ----------------
     if isinstance(checkpoint, dict):
@@ -201,8 +202,24 @@ def load_checkpoint(
             for k, v in state_dict.items()
         }
 
-    state_dict = strip_prefix(state_dict, "module.")
+    # state_dict = strip_prefix(state_dict, "module.")
+    # state_dict = strip_prefix(state_dict, "model.")
+
+    # First remove known wrappers (safe ones)
     state_dict = strip_prefix(state_dict, "model.")
+    state_dict = strip_prefix(state_dict, "_model.")
+
+    # Now handle module dynamically
+    model_keys = list(model.state_dict().keys())
+    ckpt_keys  = list(state_dict.keys())
+
+    model_has_module = model_keys[0].startswith("module.")
+    ckpt_has_module  = ckpt_keys[0].startswith("module.")
+
+    if model_has_module and not ckpt_has_module:
+        state_dict = {f"module.{k}": v for k, v in state_dict.items()}
+    elif not model_has_module and ckpt_has_module:
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
 
     # Optional custom remap
     if key_remap_fn:
@@ -229,15 +246,20 @@ def load_checkpoint(
 
     # ---------------- SCHEDULER ----------------
     if scheduler and isinstance(checkpoint, dict):
-        sch_state = checkpoint.get("scheduler") or checkpoint.get("scheduler_state_dict")
+        sch_state = checkpoint.get("scheduler") or checkpoint.get("scheduler_state_dict") or checkpoint.get("lr_scheduler")  # some use "lr_scheduler"
         if sch_state:
             scheduler.load_state_dict(sch_state)
 
+    if scaler and isinstance(checkpoint, dict):
+        scaler_state = checkpoint.get("scaler") or checkpoint.get("scaler_state_dict")
+        if scaler_state:
+            scaler.load_state_dict(scaler_state)
+
     print(f"[Checkpoint] Loaded from {ckpt_path} | epoch: {epoch}")
     print(f"Missing keys: {len(missing)}")
     print(f"Unexpected keys: {len(unexpected)}")
 
-    return model, optimizer, scheduler, epoch
+    return model, optimizer, scheduler, scaler, epoch, checkpoint
 
 
 

{opensportslib-0.1.2.dev6 → opensportslib-0.1.2.dev7/opensportslib.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opensportslib
-Version: 0.1.2.dev6
+Version: 0.1.2.dev7
 Summary: OpenSportsLib is the professional library, designed for advanced video understanding in sports. It provides state-of-the-art tools for action recognition, spotting, retrieval, and captioning, making it ideal for researchers, analysts, and developers working with sports video data.
 Author: Jeet Vora
 Requires-Python: >=3.12

{opensportslib-0.1.2.dev6 → opensportslib-0.1.2.dev7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "opensportslib"
-version = "0.1.2.dev6"
+version = "0.1.2.dev7"
 description = "OpenSportsLib is the professional library, designed for advanced video understanding in sports. It provides state-of-the-art tools for action recognition, spotting, retrieval, and captioning, making it ideal for researchers, analysts, and developers working with sports video data."
 readme = "README.md"
 requires-python = ">=3.12"