opensportslib 0.1.2.dev6__tar.gz → 0.1.2.dev8__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opensportslib
-Version: 0.1.2.dev6
+Version: 0.1.2.dev8
 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.dev8}/opensportslib/apis/base_task_model.py

@@ -9,7 +9,7 @@ import uuid
 from abc import ABC, abstractmethod
 from typing import Any
 
-from opensportslib.core.utils.config import expand, load_config_omega
+from opensportslib.core.utils.config import expand, load_config_omega, fetch_and_merge_config_from_HF
 
 
 class BaseTaskModel(ABC):
@@ -24,6 +24,13 @@ class BaseTaskModel(ABC):
         self.config_path = expand(config)
         self.config = load_config_omega(self.config_path)
 
+        if weights is not None:
+            self.config = fetch_and_merge_config_from_HF(self.config, weights, merge_policy="compatibility")
+            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)
+
         data_cfg = getattr(self.config, "DATA", None)
         if data_cfg is not None and hasattr(data_cfg, "data_dir"):
             data_cfg.data_dir = expand(data_cfg.data_dir)

{opensportslib-0.1.2.dev6 → opensportslib-0.1.2.dev8}/opensportslib/apis/classification.py

@@ -8,7 +8,6 @@ import os
 from opensportslib.apis.base_task_model import BaseTaskModel
 from opensportslib.core.utils.config import expand
 
-
 class ClassificationModel(BaseTaskModel):
     """Top-level task wrapper for classification."""
 
@@ -172,8 +171,8 @@ class ClassificationModel(BaseTaskModel):
 
         train_set = self._resolve_split_path("train", train_set)
         valid_set = self._resolve_split_path("valid", valid_set)
-
-        self.config = resolve_config_omega(self.config)
+        
+        self.config = resolve_config_omega(self.config, weights=weights)
         logging.info("Configuration:")
         logging.info(self.config)
 
@@ -241,7 +240,7 @@ class ClassificationModel(BaseTaskModel):
 
         test_set = self._resolve_split_path("test", test_set)
 
-        self.config = resolve_config_omega(self.config)
+        self.config = resolve_config_omega(self.config, weights=weights)
         logging.info("Configuration:")
         logging.info(self.config)
 
@@ -304,7 +303,7 @@ class ClassificationModel(BaseTaskModel):
 
         test_set = self._resolve_split_path("test", test_set)
 
-        self.config = resolve_config_omega(self.config)
+        self.config = resolve_config_omega(self.config, weights=weights)
         logging.info("Configuration:")
         logging.info(self.config)
         if predictions is None:

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

@@ -9,11 +9,13 @@ from opensportslib.core.utils.config import expand
 class LocalizationModel(BaseTaskModel):
     """Top-level task wrapper for localization / spotting."""
 
-    def __init__(self, config=None, weights=None):
-        super().__init__(config=config, weights=None)
-        if weights is not None:
-            self.last_loaded_weights = weights
-            self.best_checkpoint = weights
+    # def __init__(self, config=None, weights=None):
+    #     super().__init__(config=config, weights=None)
+    #     if weights is not None:
+    #         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:
@@ -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,
@@ -138,8 +178,8 @@ class LocalizationModel(BaseTaskModel):
         valid_set = self._resolve_split_path("valid", valid_set)
         self._set_split_path("train", train_set)
         self._set_split_path("valid", valid_set)
-
-        self.config = resolve_config_omega(self.config)
+        
+        self.config = resolve_config_omega(self.config, weights=weights)
         check_config(self.config, split="train")
         init_wandb(
             self.config_path,
@@ -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")
@@ -249,7 +292,7 @@ class LocalizationModel(BaseTaskModel):
         self._set_split_path("test", test_set)
 
         self.config.MODEL.multi_gpu = False
-        self.config = resolve_config_omega(self.config)
+        self.config = resolve_config_omega(self.config, weights=weights)
         check_config(self.config, split="test")
         self.config.infer_split = whether_infer_split(self.config.DATA.test)
 
@@ -318,9 +361,8 @@ class LocalizationModel(BaseTaskModel):
 
         test_set = self._resolve_split_path("test", test_set)
         self._set_split_path("test", test_set)
-
         self.config.MODEL.multi_gpu = False
-        self.config = resolve_config_omega(self.config)
+        self.config = resolve_config_omega(self.config, weights=weights)
         check_config(self.config, split="test")
         self.config.infer_split = whether_infer_split(self.config.DATA.test)
 

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

@@ -550,6 +550,13 @@ class BaseTrainerClassification:
         path_aux = os.path.join(epoch_dir, name)
         torch.save(state, path_aux)
         logging.info(f"Saved checkpoint: {path_aux}")
+
+        if self.config is not None:
+            from opensportslib.core.utils.config import save_config
+            config_path = os.path.join(epoch_dir, "config.yaml")
+            save_config(self.config, config_path)
+            logging.info(f"Saved config: {config_path}")
+            
         return path_aux
 
 
@@ -1167,11 +1174,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.dev8}/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}")
     
@@ -186,6 +155,7 @@ class Trainer_pl(Trainer):
             num_sanity_val_steps=0,
         )
         self.best_checkpoint_path = None
+        self.config = cfg
 
     def train(self, **kwargs):
         self.trainer.fit(**kwargs)
@@ -210,6 +180,13 @@ class Trainer_pl(Trainer):
             logging.info("Done training")
             logging.info(f"Best model saved at: {self.best_checkpoint_path}")
 
+            # Save the config file uniformly inside the work_dir
+            if hasattr(self, 'config') and self.config is not None:
+                from opensportslib.core.utils.config import save_config
+                config_path = os.path.join(self.work_dir, "config.yaml")
+                save_config(self.config, config_path)
+                logging.info(f"Saved config: {config_path}")
+
         log()
 
 
@@ -328,6 +305,12 @@ class Trainer_e2e(Trainer):
             self.best_checkpoint_path = best_path
             torch.save(checkpoint, best_path)
             logging.info(f"Best checkpoint saved: {best_path}")
+        
+        if self.config is not None:
+            from opensportslib.core.utils.config import save_config
+            config_path = os.path.join(self.save_dir, "config.yaml")
+            save_config(self.config, config_path)
+            logging.info(f"Saved config: {config_path}")
             
     def train(self, train_loader, valid_loader, classes):
         """Training loop with checkpoint management."""
@@ -441,7 +424,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.dev8}/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