opensportslib 0.0.1.dev2__py3-none-any.whl

opensportslib/core/sampler/weighted_sampler.py

@@ -0,0 +1,104 @@
+from torch.utils.data import DataLoader, WeightedRandomSampler
+from transformers import Trainer as HFTrainer
+import torch
+
+class WeightedTrainer(HFTrainer):
+    def __init__(self, *args, config=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.config = config
+    # --- Weighted training loader ---
+    def get_train_dataloader(self):
+        if self.train_dataset is None:
+            raise ValueError("Trainer: training requires a train_dataset.")
+        
+        weights = self.train_dataset.get_sample_weights()
+        #print(f"Sample weights: {max(weights)} {min(weights)}")
+        sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
+
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.args.per_device_train_batch_size,  # per-device batch size
+            sampler=sampler,
+            num_workers=self.args.dataloader_num_workers if hasattr(self.args, "dataloader_num_workers") else 0,
+            pin_memory=True,
+        )
+
+    def get_eval_dataloader(self, eval_dataset=None):
+        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
+        if eval_dataset is None:
+            raise ValueError("Trainer: evaluation requires an eval_dataset.")
+
+        return DataLoader(
+            eval_dataset,
+            batch_size=self.args.per_device_eval_batch_size,  # validation batch size
+            shuffle=False,  # do not shuffle evaluation
+            num_workers=self.args.dataloader_num_workers if hasattr(self.args, "dataloader_num_workers") else 0,
+            pin_memory=True,
+        )
+
+    def create_optimizer(self):
+        if self.optimizer is not None:
+            return self.optimizer
+
+        model = self.model
+
+        optimizer_grouped_parameters = []
+        print("Weighted trainer ", self.config.TRAIN.optimizer.head_lr, self.config.TRAIN.optimizer.backbone_lr)
+        # ---- Classifier head ----
+        if self.config.MODEL.unfreeze_head:
+            optimizer_grouped_parameters.append({
+                "params": model.classifier.parameters(),
+                "lr": self.config.TRAIN.optimizer.head_lr,
+            })
+
+        # ---- Backbone (last N layers) ----
+        n = self.config.MODEL.unfreeze_last_n_layers
+        if n > 0:
+            optimizer_grouped_parameters.append({
+                "params": model.videomae.encoder.layer[-n:].parameters(),
+                "lr": self.config.TRAIN.optimizer.backbone_lr,
+            })
+
+        self.optimizer = torch.optim.AdamW(
+            optimizer_grouped_parameters,
+            weight_decay=self.config.TRAIN.optimizer.weight_decay,
+        )
+
+        return self.optimizer
+
+from transformers.optimization import get_scheduler
+class VideoMAETrainer(HFTrainer):
+    def __init__(self, *args, config=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.config = config   
+
+    def create_optimizer(self):
+        if self.optimizer is not None:
+            return self.optimizer
+
+        model = self.model
+
+        optimizer_grouped_parameters = []
+        print(type(self.config.TRAIN.optimizer.head_lr), type(self.config.TRAIN.optimizer.backbone_lr))
+        # ---- Classifier head ----
+        if self.config.MODEL.unfreeze_head:
+            optimizer_grouped_parameters.append({
+                "params": model.classifier.parameters(),
+                "lr": self.config.TRAIN.optimizer.head_lr,
+            })
+
+        # ---- Backbone (last N layers) ----
+        n = self.config.MODEL.unfreeze_last_n_layers
+        if n > 0:
+            optimizer_grouped_parameters.append({
+                "params": model.videomae.encoder.layer[-n:].parameters(),
+                "lr": self.config.TRAIN.optimizer.backbone_lr,
+            })
+
+        self.optimizer = torch.optim.AdamW(
+            optimizer_grouped_parameters,
+            weight_decay=self.config.TRAIN.optimizer.weight_decay,
+        )
+
+        return self.optimizer
+

opensportslib/core/scheduler/builder.py

@@ -0,0 +1,77 @@
+"""
+Copyright 2022 James Hong, Haotian Zhang, Matthew Fisher, Michael Gharbi,
+Kayvon Fatahalian
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+this list of conditions and the following disclaimer in the documentation and/or
+other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors
+may be used to endorse or promote products derived from this software without
+specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+import torch
+from torch.optim.lr_scheduler import ChainedScheduler, LinearLR, CosineAnnealingLR
+import logging
+
+
+def build_scheduler(optimizer, cfg, default_args=None):
+    """Build a scheduler from config dict.
+
+    Args:
+        cfg (dict): Config dict. It should at least contain the key "type".
+        default_args (dict | None, optional): Default initialization arguments.
+            Default: None.
+
+    Returns:
+        scheduler: The constructed scheduler.
+    """
+    if cfg.type == "ReduceLROnPlateau":
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, mode=cfg.mode, patience=cfg.patience,
+            factor=getattr(cfg, "factor", 0.1),
+            min_lr=getattr(cfg, "min_lr", 0),
+        )
+        logging.info("Using ReduceLROnPlateau")
+    elif cfg.type == "ChainedSchedulerE2E":
+        # Warmup schedule
+        num_steps_per_epoch = default_args["len_train_loader"] // cfg.acc_grad_iter
+        cosine_epochs = cfg.num_epochs - cfg.warm_up_epochs
+        scheduler = ChainedScheduler(
+            [
+                LinearLR(
+                    optimizer,
+                    start_factor=0.01,
+                    end_factor=1.0,
+                    total_iters=cfg.warm_up_epochs * num_steps_per_epoch,
+                ),
+                CosineAnnealingLR(optimizer, num_steps_per_epoch * cosine_epochs),
+            ]
+        )
+        logging.info(
+            "Using Linear Warmup ({}) + Cosine Annealing LR ({})".format(
+                cfg.warm_up_epochs, cosine_epochs
+            )
+        )
+    elif cfg.type == "StepLR":
+        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=cfg.step_size, gamma=cfg.gamma)
+    else:
+        scheduler = None
+    return scheduler