opensportslib 0.1.2.dev5__tar.gz → 0.1.2.dev7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opensportslib
-Version: 0.1.2.dev5
+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.dev5 → 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,9 +207,6 @@ class LocalizationModel(BaseTaskModel):
 
         start = time.time()
 
-        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,
@@ -196,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.dev5 → 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.dev5 → 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.dev5 → 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.dev5 → opensportslib-0.1.2.dev7}/opensportslib/datasets/localization_dataset.py

@@ -1016,7 +1016,10 @@ if DALI_AVAILABLE:
             for pipe in self.pipes:
                 pipe.build()
 
-            super().__init__(self.pipes, output_map, size=self.nb_videos)
+            # Pipeline returns (video, label_idx, frame_num) - label processing
+            # is done post-hoc in get_attr to avoid DALI 2.0 fn.python_function issues
+            internal_output_map = ['data', 'label_idx', 'frame_num']
+            super().__init__(self.pipes, internal_output_map, size=self.nb_videos)
 
             self.device = torch.device(
                 "cuda:{}".format(self.devices[1 if len(self.devices) > 1 else 0])
@@ -1052,8 +1055,19 @@ if DALI_AVAILABLE:
             Returns:
                 dict :{"frames","contains_event","labels"}.
             """
-            batch_labels = batch["label"]
+            batch_label_idx = batch["label_idx"]
+            batch_frame_num = batch["frame_num"]
             batch_images = batch["data"]
+
+            batch_size = batch_label_idx.shape[0]
+            batch_labels = torch.zeros(batch_size, self.clip_len, dtype=torch.int64)
+            for b in range(batch_size):
+                video_idx = int(batch_label_idx[b].item())
+                frame_num = int(batch_frame_num[b].item())
+                batch_labels[b] = torch.from_numpy(
+                    self._compute_labels(video_idx, frame_num)
+                )
+
             sum_labels = torch.sum(
                 batch_labels, dim=1 if len(batch_labels.shape) == 2 else 0
             )
@@ -1229,26 +1243,22 @@ if DALI_AVAILABLE:
                     std=[255, 255, 255],
                     mirror=fn.random.coin_flip(),
                 )
-            label = fn.python_function(
-                label, frame_num, function=self.edit_labels, device="gpu"
-            )
-            return video, label
+            return video, label, frame_num
 
-        def edit_labels(self, label, frame_num):
-            """Construct a list having the same length as the number of frames. The elements of the list are the indexes (starting at 1) of the class where an event occurs, 0 otherwise.
+        def _compute_labels(self, video_idx, frame_num):
+            """Construct a label array for a clip. Each element is the class index
+            (starting at 1) where an event occurs, 0 otherwise.
 
             Args:
-                label :index of the video to get the metadata.
-                frame_num :index of start frame.
+                video_idx (int): Index of the video in self._labels.
+                frame_num (int): Raw start frame number from the reader.
 
             Returns:
-                labels (cupy.array): the list of labels (corresponding to events) corresponding with the extracted frames.
+                labels (np.ndarray): Label array of shape (clip_len,).
             """
-            import cupy
-
-            video_meta = self._labels[label.item()]
-            base_idx = frame_num.item() // self._stride
-            labels = cupy.zeros(self.clip_len, np.int64)
+            video_meta = self._labels[video_idx]
+            base_idx = frame_num // self._stride
+            labels = np.zeros(self.clip_len, np.int64)
 
             for event in video_meta["events"]:
                 event_frame = event["frame"]
@@ -1258,12 +1268,12 @@ if DALI_AVAILABLE:
                     label_idx >= self.dilate_len
                     and label_idx < self.clip_len + self.dilate_len
                 ):
-                    label = self._class_dict[event["label"]]
+                    label_val = self._class_dict[event["label"]]
                     for i in range(
                         max(0, label_idx - self.dilate_len),
                         min(self.clip_len, label_idx + self.dilate_len + 1),
                     ):
-                        labels[i] = label
+                        labels[i] = label_val
             return labels
 
         def print_info(self):

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opensportslib
-Version: 0.1.2.dev5
+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.dev5 → opensportslib-0.1.2.dev7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "opensportslib"
-version = "0.1.2.dev5"
+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"