sleap-nn 0.1.0__py3-none-any.whl → 0.1.0a1__py3-none-any.whl

sleap_nn/training/schedulers.py

@@ -1,191 +0,0 @@
-"""Custom learning rate schedulers for sleap-nn training.
-
-This module provides learning rate schedulers with warmup phases that are commonly
-used in deep learning for pose estimation and computer vision tasks.
-"""
-
-import math
-from torch.optim.lr_scheduler import LRScheduler
-
-
-class LinearWarmupCosineAnnealingLR(LRScheduler):
-    """Cosine annealing learning rate scheduler with linear warmup.
-
-    The learning rate increases linearly from `warmup_start_lr` to the optimizer's
-    base learning rate over `warmup_epochs`, then decreases following a cosine
-    curve to `eta_min` over the remaining epochs.
-
-    This schedule is widely used in vision transformers and modern CNN architectures
-    as it provides stable early training (warmup) and smooth convergence (cosine decay).
-
-    Args:
-        optimizer: Wrapped optimizer.
-        warmup_epochs: Number of epochs for the linear warmup phase.
-        max_epochs: Total number of training epochs.
-        warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
-        eta_min: Minimum learning rate at the end of the schedule. Default: 0.0.
-        last_epoch: The index of the last epoch. Default: -1.
-
-    Example:
-        >>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
-        >>> scheduler = LinearWarmupCosineAnnealingLR(
-        ...     optimizer, warmup_epochs=5, max_epochs=100, eta_min=1e-6
-        ... )
-        >>> for epoch in range(100):
-        ...     train(...)
-        ...     scheduler.step()
-    """
-
-    def __init__(
-        self,
-        optimizer,
-        warmup_epochs: int,
-        max_epochs: int,
-        warmup_start_lr: float = 0.0,
-        eta_min: float = 0.0,
-        last_epoch: int = -1,
-    ):
-        """Initialize the scheduler.
-
-        Args:
-            optimizer: Wrapped optimizer.
-            warmup_epochs: Number of epochs for the linear warmup phase.
-            max_epochs: Total number of training epochs.
-            warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
-            eta_min: Minimum learning rate at the end of the schedule. Default: 0.0.
-            last_epoch: The index of the last epoch. Default: -1.
-        """
-        if warmup_epochs < 0:
-            raise ValueError(f"warmup_epochs must be >= 0, got {warmup_epochs}")
-        if max_epochs <= 0:
-            raise ValueError(f"max_epochs must be > 0, got {max_epochs}")
-        if warmup_epochs >= max_epochs:
-            raise ValueError(
-                f"warmup_epochs ({warmup_epochs}) must be < max_epochs ({max_epochs})"
-            )
-        if warmup_start_lr < 0:
-            raise ValueError(f"warmup_start_lr must be >= 0, got {warmup_start_lr}")
-        if eta_min < 0:
-            raise ValueError(f"eta_min must be >= 0, got {eta_min}")
-
-        self.warmup_epochs = warmup_epochs
-        self.max_epochs = max_epochs
-        self.warmup_start_lr = warmup_start_lr
-        self.eta_min = eta_min
-        super().__init__(optimizer, last_epoch)
-
-    def get_lr(self):
-        """Compute the learning rate at the current epoch."""
-        if self.last_epoch < self.warmup_epochs:
-            # Linear warmup phase
-            if self.warmup_epochs == 0:
-                return list(self.base_lrs)
-            alpha = self.last_epoch / self.warmup_epochs
-            return [
-                self.warmup_start_lr + alpha * (base_lr - self.warmup_start_lr)
-                for base_lr in self.base_lrs
-            ]
-        else:
-            # Cosine annealing phase
-            decay_epochs = self.max_epochs - self.warmup_epochs
-            if decay_epochs == 0:
-                return [self.eta_min for _ in self.base_lrs]
-            progress = (self.last_epoch - self.warmup_epochs) / decay_epochs
-            # Clamp progress to [0, 1] to handle epochs beyond max_epochs
-            progress = min(1.0, progress)
-            return [
-                self.eta_min
-                + (base_lr - self.eta_min) * (1 + math.cos(math.pi * progress)) / 2
-                for base_lr in self.base_lrs
-            ]
-
-
-class LinearWarmupLinearDecayLR(LRScheduler):
-    """Linear warmup followed by linear decay learning rate scheduler.
-
-    The learning rate increases linearly from `warmup_start_lr` to the optimizer's
-    base learning rate over `warmup_epochs`, then decreases linearly to `end_lr`
-    over the remaining epochs.
-
-    This schedule provides a simple, interpretable learning rate trajectory and is
-    commonly used in transformer-based models and NLP tasks.
-
-    Args:
-        optimizer: Wrapped optimizer.
-        warmup_epochs: Number of epochs for the linear warmup phase.
-        max_epochs: Total number of training epochs.
-        warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
-        end_lr: Learning rate at the end of training. Default: 0.0.
-        last_epoch: The index of the last epoch. Default: -1.
-
-    Example:
-        >>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
-        >>> scheduler = LinearWarmupLinearDecayLR(
-        ...     optimizer, warmup_epochs=5, max_epochs=100, end_lr=1e-6
-        ... )
-        >>> for epoch in range(100):
-        ...     train(...)
-        ...     scheduler.step()
-    """
-
-    def __init__(
-        self,
-        optimizer,
-        warmup_epochs: int,
-        max_epochs: int,
-        warmup_start_lr: float = 0.0,
-        end_lr: float = 0.0,
-        last_epoch: int = -1,
-    ):
-        """Initialize the scheduler.
-
-        Args:
-            optimizer: Wrapped optimizer.
-            warmup_epochs: Number of epochs for the linear warmup phase.
-            max_epochs: Total number of training epochs.
-            warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
-            end_lr: Learning rate at the end of training. Default: 0.0.
-            last_epoch: The index of the last epoch. Default: -1.
-        """
-        if warmup_epochs < 0:
-            raise ValueError(f"warmup_epochs must be >= 0, got {warmup_epochs}")
-        if max_epochs <= 0:
-            raise ValueError(f"max_epochs must be > 0, got {max_epochs}")
-        if warmup_epochs >= max_epochs:
-            raise ValueError(
-                f"warmup_epochs ({warmup_epochs}) must be < max_epochs ({max_epochs})"
-            )
-        if warmup_start_lr < 0:
-            raise ValueError(f"warmup_start_lr must be >= 0, got {warmup_start_lr}")
-        if end_lr < 0:
-            raise ValueError(f"end_lr must be >= 0, got {end_lr}")
-
-        self.warmup_epochs = warmup_epochs
-        self.max_epochs = max_epochs
-        self.warmup_start_lr = warmup_start_lr
-        self.end_lr = end_lr
-        super().__init__(optimizer, last_epoch)
-
-    def get_lr(self):
-        """Compute the learning rate at the current epoch."""
-        if self.last_epoch < self.warmup_epochs:
-            # Linear warmup phase
-            if self.warmup_epochs == 0:
-                return list(self.base_lrs)
-            alpha = self.last_epoch / self.warmup_epochs
-            return [
-                self.warmup_start_lr + alpha * (base_lr - self.warmup_start_lr)
-                for base_lr in self.base_lrs
-            ]
-        else:
-            # Linear decay phase
-            decay_epochs = self.max_epochs - self.warmup_epochs
-            if decay_epochs == 0:
-                return [self.end_lr for _ in self.base_lrs]
-            progress = (self.last_epoch - self.warmup_epochs) / decay_epochs
-            # Clamp progress to [0, 1] to handle epochs beyond max_epochs
-            progress = min(1.0, progress)
-            return [
-                base_lr + progress * (self.end_lr - base_lr)
-                for base_lr in self.base_lrs
-            ]

sleap_nn-0.1.0.dist-info/RECORD

@@ -1,88 +0,0 @@
-sleap_nn/.DS_Store,sha256=HY8amA79eHkt7o5VUiNsMxkc9YwW6WIPyZbYRj_JdSU,6148
-sleap_nn/__init__.py,sha256=nl_xOwr4X0BDdND9odwKF63zX2EytLkDx0S4EDMJjeo,1360
-sleap_nn/cli.py,sha256=VU0fHIy02BErq8aQs0s9bNAKx__qpwzbfpsA7j3jHVI,31285
-sleap_nn/evaluation.py,sha256=SRO3qNOyyGoNBLLA2OKIUhvwyk0oI2ax1rtYmccx6m0,33785
-sleap_nn/legacy_models.py,sha256=fKUzmsmHpkEACk2uSDYcILoV1gNFoOXwxlzP_NO3jWo,22426
-sleap_nn/predict.py,sha256=ai-0nfknh6JT0BrJL_DECq6QVRe1cCwfpvXQCaDVC7o,36515
-sleap_nn/system_info.py,sha256=7tWe3y6s872nDbrZoHIdSs-w4w46Z4dEV2qCV-Fe7No,14711
-sleap_nn/train.py,sha256=afRzm5t0S1GTFXAkPB1Qh4sa71sIwnNwFZjmvZ33wwI,30765
-sleap_nn/architectures/__init__.py,sha256=w0XxQcx-CYyooszzvxRkKWiJkUg-26IlwQoGna8gn40,46
-sleap_nn/architectures/common.py,sha256=MLv-zdHsWL5Q2ct_Wv6SQbRS-5hrFtjK_pvBEfwx-vU,3660
-sleap_nn/architectures/convnext.py,sha256=Ba9SFQHBdfz8gcMYZPMItuW-FyQuHBgUU0M8MWhaHuY,14210
-sleap_nn/architectures/encoder_decoder.py,sha256=1cBk9WU0zkXC2aK9XZy6VKHEe2hJEpIa-rwCxNgObZg,29292
-sleap_nn/architectures/heads.py,sha256=5E-7kQ-b2gsL0EviQ8z3KS1DAAMT4F2ZnEzx7eSG5gg,21001
-sleap_nn/architectures/model.py,sha256=1_dsP_4T9fsEVJjDt3er0haMKtbeM6w6JC6tc2jD0Gw,7139
-sleap_nn/architectures/swint.py,sha256=hlShh1Br0eTijir2U3np8sAaNJa12Xny0VzPx8HSaRo,15060
-sleap_nn/architectures/unet.py,sha256=rAy2Omi6tv1MNW2nBn0Tw-94Nw_-1wFfCT3-IUyPcgo,11723
-sleap_nn/architectures/utils.py,sha256=L0KVs0gbtG8U75Sl40oH_r_w2ySawh3oQPqIGi54HGo,2171
-sleap_nn/config/__init__.py,sha256=l0xV1uJsGJfMPfWAqlUR7Ivu4cSCWsP-3Y9ueyPESuk,42
-sleap_nn/config/data_config.py,sha256=LiZ3TGUZk6EiLJXXsu9kGOf84eUyGNMWrQ3XvsIYuAA,24426
-sleap_nn/config/get_config.py,sha256=er-ZTMQT68uhDNbQUDVhY8FzqeyFDcgEpLn1ZRMjATQ,42449
-sleap_nn/config/model_config.py,sha256=XFIbqFno7IkX0Se5WF_2_7aUalAlC2SvpDe-uP2TttM,57582
-sleap_nn/config/trainer_config.py,sha256=JBbojDcxDZR5RqpM4eNqEko6-nOYuvT3eDPSpb1OLFQ,31715
-sleap_nn/config/training_job_config.py,sha256=v12_ME_tBUg8JFwOxJNW4sDQn-SedDhiJOGz-TlRwT0,5861
-sleap_nn/config/utils.py,sha256=GgWgVs7_N7ifsJ5OQG3_EyOagNyN3Dx7wS2BAlkaRkg,5553
-sleap_nn/data/__init__.py,sha256=eMNvFJFa3gv5Rq8oK5wzo6zt1pOlwUGYf8EQii6bq7c,54
-sleap_nn/data/augmentation.py,sha256=1tTHphaD021SYvt27dpGnlgZAxjRIcGetnWvz20YXtc,6454
-sleap_nn/data/confidence_maps.py,sha256=PTRqZWSAz1S7viJhxu7QgIC1aHiek97c_dCUsKUwG1o,6217
-sleap_nn/data/custom_datasets.py,sha256=CdZ_hblKHvIf4zYD0CWGU511uA3_7Ca1oq2r3mbyz5I,112101
-sleap_nn/data/edge_maps.py,sha256=75qG_7zHRw7fC8JUCVI2tzYakIoxxneWWmcrTwjcHPo,12519
-sleap_nn/data/identity.py,sha256=7vNup6PudST4yDLyDT9wDO-cunRirTEvx4sP77xrlfk,5193
-sleap_nn/data/instance_centroids.py,sha256=SF-3zJt_VMTbZI5ssbrvmZQZDd3684bn55EAtvcbQ6o,2172
-sleap_nn/data/instance_cropping.py,sha256=1DQJDmSM18gUVofXfoVg-hQnXPeLdtY-buGSbTPLy1c,8334
-sleap_nn/data/normalization.py,sha256=5xEvcguG-fvAGObl4nWPZ9TEM5gvv0uYPGDuni34XII,2930
-sleap_nn/data/providers.py,sha256=0x6GFP1s1c08ji4p0M5V6p-dhT4Z9c-SI_Aw1DWX-uM,14272
-sleap_nn/data/resizing.py,sha256=bYxELbDJZen3tiwEbv__EtNgmJcQJhQ7CRIJ30t-UvI,5040
-sleap_nn/data/skia_augmentation.py,sha256=CQYy592-poPfiwOcTR6X9eY1lm-e7piJU-LlBHEe37w,14593
-sleap_nn/data/utils.py,sha256=ap3fYqk2UFiCS4cAdZz8x3ahMrkLmvXnHfRx6R-_5OQ,11130
-sleap_nn/export/__init__.py,sha256=E5FN_dG7583V_poTMS8jthL0Bf07qFCUA0kY2iewlvk,574
-sleap_nn/export/cli.py,sha256=UcXrefi7YKunaIzbW97QeBQmbG3_-7Xa4XZgcwOGd9k,63568
-sleap_nn/export/metadata.py,sha256=ory-smt2tpgObmld7WicxEFFHM03hl6yasl4N3KVTmA,7427
-sleap_nn/export/utils.py,sha256=oSVf_PCG1IUZd6ZpgwOMPgqVkMxaO50tt2C-cTUMGg4,11414
-sleap_nn/export/exporters/__init__.py,sha256=zBZMEjqzjGDq1KvIiKSVrHPplE4zVtUW3MHQ5kGg1M4,1459
-sleap_nn/export/exporters/onnx_exporter.py,sha256=i3GWYn9NH2i_SJd2OmVHZgUfJJW9w-Swms1YDM-HhEc,2129
-sleap_nn/export/exporters/tensorrt_exporter.py,sha256=5v_wnkal7bPKgrhQRL8akYikmQ_8PdqVr1ajpmU1oGo,9069
-sleap_nn/export/predictors/__init__.py,sha256=xmDFeK05eHO9-Obue3jhBy4VMBW0NZjyGXSuL3RMhJY,1772
-sleap_nn/export/predictors/base.py,sha256=5uJjMSuG4p-I3YPY9tpn6m8HdENcPqFtjpehL5mNuAg,584
-sleap_nn/export/predictors/onnx.py,sha256=uiFIrCQBxZ_pK4x0JUx2N2k6qzN8KqfQtXAa4FkqVIs,5125
-sleap_nn/export/predictors/tensorrt.py,sha256=lrYjw_lasjbGIku8aM5b26G92bXdvQaVAKrT6U4Xcco,10805
-sleap_nn/export/wrappers/__init__.py,sha256=OeMCP8UCXNwmAYCrudKX-Gt4nhr9z13bzkdrZHjUNnk,984
-sleap_nn/export/wrappers/base.py,sha256=sd55SCkjOdMYut9Yh8Yi1MMF8EY0U9oYtXadmi7xD80,3462
-sleap_nn/export/wrappers/bottomup.py,sha256=rh4-wu4Oc8JWWW3K3xLhiTkyQ3dqhJ9QfMrPY65ck4Y,9237
-sleap_nn/export/wrappers/bottomup_multiclass.py,sha256=r0-LJxv8zYTUr94ytlj1rp0GDuodY6CugLIl93IVhRA,7495
-sleap_nn/export/wrappers/centered_instance.py,sha256=PShtVAlJhSbyuAVmRNjIM3PnOngaRFX3LtX-wO91Qqs,1704
-sleap_nn/export/wrappers/centroid.py,sha256=3nbbysCkzHS9pRC2-37z4Z1Ol5XEdvYahyR4Gm25dqc,1840
-sleap_nn/export/wrappers/single_instance.py,sha256=ZoQICK5_tuPjA19wsDmq8ljhM1Qus2_O8ZCDq6yd2L4,2826
-sleap_nn/export/wrappers/topdown.py,sha256=_HAtEASBSfLTfFO9Q7honIpm4ECNoMR76JLkWyO-PF8,6809
-sleap_nn/export/wrappers/topdown_multiclass.py,sha256=fJ3SC76BdB5QWOvMobCoKqf1OQBJGoi5MS0Bv6_gQmY,11920
-sleap_nn/inference/__init__.py,sha256=eVkCmKrxHlDFJIlZTf8B5XEOcSyw-gPQymXMY5uShOM,170
-sleap_nn/inference/bottomup.py,sha256=3s90aRlpIcRnSNe-R5-qiuX3S48kCWMpCl8YuNnTEDI,17084
-sleap_nn/inference/identity.py,sha256=GjNDL9MfGqNyQaK4AE8JQCAE8gpMuE_Y-3r3Gpa53CE,6540
-sleap_nn/inference/paf_grouping.py,sha256=7Fo9lCAj-zcHgv5rI5LIMYGcixCGNt_ZbSNs8Dik7l8,69973
-sleap_nn/inference/peak_finding.py,sha256=JS4qIcl72Ep4vbbYYdngqu7rRl3MW8kFaH8W9BRoOyA,15920
-sleap_nn/inference/postprocessing.py,sha256=ZM_OH7_WIprieaujZ2Rk_34JhSDDzCry6Pq2YM_u5sg,8998
-sleap_nn/inference/predictors.py,sha256=dztVxdEPgSrc8bfB-t4DBa0wkg9p0HcVoZrZ-FcpP24,160720
-sleap_nn/inference/provenance.py,sha256=0BekXyvpLMb0Vv6DjpctlLduG9RN-Q8jt5zDm783eZE,11204
-sleap_nn/inference/single_instance.py,sha256=rOns_5TsJ1rb-lwmHG3ZY-pOhXGN2D-SfW9RmBxxzcI,4089
-sleap_nn/inference/topdown.py,sha256=Ha0Nwx-XCH_rebIuIGhP0qW68QpjLB3XRr9rxt05JLs,35108
-sleap_nn/inference/utils.py,sha256=JnaJK4S_qLtHkWOSkHf4oRZjOmgnU9BGADQnntgGxxs,4689
-sleap_nn/tracking/__init__.py,sha256=rGR35wpSW-n5d3cMiQUzQQ_Dy5II5DPjlXAoPw2QhmM,31
-sleap_nn/tracking/track_instance.py,sha256=9k0uVy9VmpleaLcJh7sVWSeFUPXiw7yj95EYNdXJcks,1373
-sleap_nn/tracking/tracker.py,sha256=_WT-HFruzyOsvcq3AtLm3vnI9MYSwyBmq-HlQvj1vmU,41955
-sleap_nn/tracking/utils.py,sha256=uHVd_mzzZjviVDdLSKXJJ1T96n5ObKvkqIuGsl9Yy8U,11276
-sleap_nn/tracking/candidates/__init__.py,sha256=1O7NObIwshM7j1rLHmImbFphvkM9wY1j4j1TvO5scSE,49
-sleap_nn/tracking/candidates/fixed_window.py,sha256=D80KMlTnenuQveQVVhk9j0G8yx6K324C7nMLHgG76e0,6296
-sleap_nn/tracking/candidates/local_queues.py,sha256=Nx3R5wwEwq0gbfH-fi3oOumfkQo8_sYe5GN47pD9Be8,7305
-sleap_nn/training/__init__.py,sha256=vNTKsIJPZHJwFSKn5PmjiiRJunR_9e7y4_v0S6rdF8U,32
-sleap_nn/training/callbacks.py,sha256=k1WAe7ZdanwPyrJua08-wrzoD-fjWxPk252Q4ncL-Gs,63250
-sleap_nn/training/lightning_modules.py,sha256=Py6aFV51URL-SOqjO-cbdS4VsjNoNwMj0RcAFAqQe_E,110241
-sleap_nn/training/losses.py,sha256=gbdinUURh4QUzjmNd2UJpt4FXwecqKy9gHr65JZ1bZk,1632
-sleap_nn/training/model_trainer.py,sha256=iYm2jGhNXNt6jc-FZRaIS6kJVOvdX0Ixk7lnvQEZBLc,59844
-sleap_nn/training/schedulers.py,sha256=PFc1M-PM1khaNfndoi5VSW-A04lLYfPf_5gFA4pkrLk,7681
-sleap_nn/training/utils.py,sha256=ivdkZEI0DkTCm6NPszsaDOh9jSfozkONZdl6TvvQUWI,20398
-sleap_nn-0.1.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-sleap_nn-0.1.0.dist-info/METADATA,sha256=ta-irf2lxIi8zfDo8kb-nF_AobfqEX0_vsVS-5tQDzU,6302
-sleap_nn-0.1.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-sleap_nn-0.1.0.dist-info/entry_points.txt,sha256=zfl5Y3hidZxWBvo8qXvu5piJAXJ_l6v7xVFm0gNiUoI,46
-sleap_nn-0.1.0.dist-info/top_level.txt,sha256=Kz68iQ55K75LWgSeqz4V4SCMGeFFYH-KGBOyhQh3xZE,9
-sleap_nn-0.1.0.dist-info/RECORD,,