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

sleap_nn/training/model_trainer.py

@@ -23,7 +23,6 @@ from sleap_nn.data.utils import check_cache_memory
 from lightning.pytorch.callbacks import (
     ModelCheckpoint,
     EarlyStopping,
-    LearningRateMonitor,
 )
 from lightning.pytorch.profilers import (
     SimpleProfiler,
@@ -61,6 +60,7 @@ from sleap_nn.training.callbacks import (
     WandBVizCallbackWithPAFs,
     CSVLoggerCallback,
     SleapProgressBar,
+    EpochEndEvaluationCallback,
 )
 from sleap_nn import RANK
 from sleap_nn.legacy_models import get_keras_first_layer_channels
@@ -493,13 +493,10 @@ class ModelTrainer:
         if run_name is None or run_name == "" or run_name == "None":
             sum_train_lfs = sum([len(train_label) for train_label in self.train_labels])
             sum_val_lfs = sum([len(val_label) for val_label in self.val_labels])
-            if self._get_trainer_devices() > 1:
-                run_name = f"{self.model_type}.n={sum_train_lfs + sum_val_lfs}"
-            else:
-                run_name = (
-                    datetime.now().strftime("%y%m%d_%H%M%S")
-                    + f".{self.model_type}.n={sum_train_lfs + sum_val_lfs}"
-                )
+            run_name = (
+                datetime.now().strftime("%y%m%d_%H%M%S")
+                + f".{self.model_type}.n={sum_train_lfs + sum_val_lfs}"
+            )
 
         # If checkpoint path already exists, add suffix to prevent overwriting
         if (Path(ckpt_dir) / run_name).exists() and (
@@ -642,10 +639,6 @@ class ModelTrainer:
         if self.config.trainer_config.wandb.prv_runid == "":
             self.config.trainer_config.wandb.prv_runid = None
 
-        # Default wandb run name to trainer run_name if not specified
-        if self.config.trainer_config.wandb.name is None:
-            self.config.trainer_config.wandb.name = self.config.trainer_config.run_name
-
         # compute preprocessing parameters from the labels objects and fill in the config
         self._setup_preprocessing_config()
 
@@ -695,9 +688,14 @@ class ModelTrainer:
                 )
             )
 
-        # setup checkpoint path
+        # setup checkpoint path (generates run_name if not specified)
         self._setup_ckpt_path()
 
+        # Default wandb run name to trainer run_name if not specified
+        # Note: This must come after _setup_ckpt_path() which generates run_name
+        if self.config.trainer_config.wandb.name is None:
+            self.config.trainer_config.wandb.name = self.config.trainer_config.run_name
+
         # verify input_channels in model_config based on input image and pretrained model weights
         self._verify_model_input_channels()
 
@@ -837,7 +835,7 @@ class ModelTrainer:
                     / self.config.trainer_config.run_name
                 ).as_posix(),
                 filename="best",
-                monitor="val_loss",
+                monitor="val/loss",
                 mode="min",
             )
             callbacks.append(checkpoint_callback)
@@ -845,18 +843,20 @@ class ModelTrainer:
             # csv log callback
             csv_log_keys = [
                 "epoch",
-                "train_loss",
-                "val_loss",
+                "train/loss",
+                "val/loss",
                 "learning_rate",
-                "train_time",
-                "val_time",
+                "train/time",
+                "val/time",
             ]
             if self.model_type in [
                 "single_instance",
                 "centered_instance",
                 "multi_class_topdown",
             ]:
-                csv_log_keys.extend(self.skeletons[0].node_names)
+                csv_log_keys.extend(
+                    [f"train/confmaps/{name}" for name in self.skeletons[0].node_names]
+                )
             csv_logger = CSVLoggerCallback(
                 filepath=Path(self.config.trainer_config.ckpt_dir)
                 / self.config.trainer_config.run_name
@@ -869,7 +869,7 @@ class ModelTrainer:
             # early stopping callback
             callbacks.append(
                 EarlyStopping(
-                    monitor="val_loss",
+                    monitor="val/loss",
                     mode="min",
                     verbose=False,
                     min_delta=self.config.trainer_config.early_stopping.min_delta,
@@ -909,10 +909,6 @@ class ModelTrainer:
                     "To keep logs, set trainer_config.wandb.delete_local_logs=false"
                 )
 
-            # Learning rate monitor callback - logs LR at each step for dynamic schedulers
-            # Only added when wandb is enabled since it requires a logger
-            callbacks.append(LearningRateMonitor(logging_interval="step"))
-
             # save the configs as yaml in the checkpoint dir
             # Mask API key in both configs to prevent saving to disk
             self.config.trainer_config.wandb.api_key = ""
@@ -1086,6 +1082,18 @@ class ModelTrainer:
         if self.config.trainer_config.enable_progress_bar:
             callbacks.append(SleapProgressBar())
 
+        # Add epoch-end evaluation callback if enabled
+        if self.config.trainer_config.eval.enabled:
+            callbacks.append(
+                EpochEndEvaluationCallback(
+                    skeleton=self.skeletons[0],
+                    videos=self.val_labels[0].videos,
+                    eval_frequency=self.config.trainer_config.eval.frequency,
+                    oks_stddev=self.config.trainer_config.eval.oks_stddev,
+                    oks_scale=self.config.trainer_config.eval.oks_scale,
+                )
+            )
+
         return loggers, callbacks
 
     def _delete_cache_imgs(self):
@@ -1268,18 +1276,21 @@ class ModelTrainer:
                 # Define custom x-axes for wandb metrics
                 # Epoch-level metrics use epoch as x-axis, step-level use default global_step
                 wandb.define_metric("epoch")
-                wandb.define_metric("val_loss", step_metric="epoch")
-                wandb.define_metric("val_time", step_metric="epoch")
-                wandb.define_metric("train_time", step_metric="epoch")
-                # Per-node losses use epoch as x-axis
-                for node_name in self.skeletons[0].node_names:
-                    wandb.define_metric(node_name, step_metric="epoch")
-
-                # Visualization images use epoch as x-axis
-                wandb.define_metric("train_predictions*", step_metric="epoch")
-                wandb.define_metric("val_predictions*", step_metric="epoch")
-                wandb.define_metric("train_pafs*", step_metric="epoch")
-                wandb.define_metric("val_pafs*", step_metric="epoch")
+
+                # Training metrics (train/ prefix for grouping) - all use epoch x-axis
+                wandb.define_metric("train/*", step_metric="epoch")
+                wandb.define_metric("train/confmaps/*", step_metric="epoch")
+
+                # Validation metrics (val/ prefix for grouping)
+                wandb.define_metric("val/*", step_metric="epoch")
+
+                # Evaluation metrics (eval/ prefix for grouping)
+                wandb.define_metric("eval/*", step_metric="epoch")
+
+                # Visualization images (need explicit nested paths)
+                wandb.define_metric("viz/*", step_metric="epoch")
+                wandb.define_metric("viz/train/*", step_metric="epoch")
+                wandb.define_metric("viz/val/*", step_metric="epoch")
 
                 self.config.trainer_config.wandb.current_run_id = wandb.run.id
                 wandb.config["run_name"] = self.config.trainer_config.wandb.name
@@ -1322,27 +1333,7 @@ class ModelTrainer:
             logger.info(
                 f"Finished training loop. [{(time.time() - start_train_time) / 60:.1f} min]"
             )
-            if self.trainer.global_rank == 0 and self.config.trainer_config.use_wandb:
-                wandb.finish()
-
-                # Delete local wandb logs if configured
-                wandb_config = self.config.trainer_config.wandb
-                should_delete_wandb_logs = wandb_config.delete_local_logs is True or (
-                    wandb_config.delete_local_logs is None
-                    and wandb_config.wandb_mode != "offline"
-                )
-                if should_delete_wandb_logs:
-                    wandb_dir = (
-                        Path(self.config.trainer_config.ckpt_dir)
-                        / self.config.trainer_config.run_name
-                        / "wandb"
-                    )
-                    if wandb_dir.exists():
-                        logger.info(
-                            f"Deleting local wandb logs at {wandb_dir}... "
-                            "(set trainer_config.wandb.delete_local_logs=false to disable)"
-                        )
-                        shutil.rmtree(wandb_dir, ignore_errors=True)
+            # Note: wandb.finish() is called in train.py after post-training evaluation
 
             # delete image disk caching
             if (

{sleap_nn-0.1.0a1.dist-info → sleap_nn-0.1.0a3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleap-nn
-Version: 0.1.0a1
+Version: 0.1.0a3
 Summary: Neural network backend for training and inference for animal pose estimation.
 Author-email: Divya Seshadri Murali <dimurali@salk.edu>, Elizabeth Berrigan <eberrigan@salk.edu>, Vincent Tu <vitu@ucsd.edu>, Liezl Maree <lmaree@salk.edu>, David Samy <davidasamy@gmail.com>, Talmo Pereira <talmo@salk.edu>
 License: BSD-3-Clause
@@ -13,7 +13,7 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: sleap-io<0.7.0,>=0.6.0
+Requires-Dist: sleap-io<0.7.0,>=0.6.2
 Requires-Dist: numpy
 Requires-Dist: lightning
 Requires-Dist: kornia
@@ -47,6 +47,17 @@ Requires-Dist: torchvision>=0.20.0; extra == "torch-cuda128"
 Provides-Extra: torch-cuda130
 Requires-Dist: torch; extra == "torch-cuda130"
 Requires-Dist: torchvision>=0.20.0; extra == "torch-cuda130"
+Provides-Extra: export
+Requires-Dist: onnx>=1.15.0; extra == "export"
+Requires-Dist: onnxruntime>=1.16.0; extra == "export"
+Requires-Dist: onnxscript>=0.1.0; extra == "export"
+Provides-Extra: export-gpu
+Requires-Dist: onnx>=1.15.0; extra == "export-gpu"
+Requires-Dist: onnxruntime-gpu>=1.16.0; extra == "export-gpu"
+Requires-Dist: onnxscript>=0.1.0; extra == "export-gpu"
+Provides-Extra: tensorrt
+Requires-Dist: tensorrt>=10.0.0; extra == "tensorrt"
+Requires-Dist: torch-tensorrt>=2.0.0; extra == "tensorrt"
 Dynamic: license-file
 
 # sleap-nn

{sleap_nn-0.1.0a1.dist-info → sleap_nn-0.1.0a3.dist-info}/RECORD

@@ -1,11 +1,11 @@
 sleap_nn/.DS_Store,sha256=HY8amA79eHkt7o5VUiNsMxkc9YwW6WIPyZbYRj_JdSU,6148
-sleap_nn/__init__.py,sha256=l5Lwiad8GOurqkAhMwWw8-UcpH6af2TnMURf-oKj_U8,1362
-sleap_nn/cli.py,sha256=U4hpEcOxK7a92GeItY95E2DRm5P1ME1GqU__mxaDcW0,21167
-sleap_nn/evaluation.py,sha256=3u7y85wFoBgCwOB2xOGTJIDrd2dUPWOo4m0s0oW3da4,31095
+sleap_nn/__init__.py,sha256=rBA8WZU2_28NOBxfsqkSeA1hMG4pKEPiB2_yZYmhCIg,1362
+sleap_nn/cli.py,sha256=O_edfRK6imWE8KFECQ-Ux0sqlLLbsyQYQjTLtQ63ry4,22295
+sleap_nn/evaluation.py,sha256=SRO3qNOyyGoNBLLA2OKIUhvwyk0oI2ax1rtYmccx6m0,33785
 sleap_nn/legacy_models.py,sha256=8aGK30DZv3pW2IKDBEWH1G2mrytjaxPQD4miPUehj0M,20258
-sleap_nn/predict.py,sha256=8QKjRbS-L-6HF1NFJWioBPv3HSzUpFr2oGEB5hRJzQA,35523
+sleap_nn/predict.py,sha256=B0YhCdJ35TLOydpVvc3gHdkN45pf_ruhi7jZ_6Rrfqs,37089
 sleap_nn/system_info.py,sha256=7tWe3y6s872nDbrZoHIdSs-w4w46Z4dEV2qCV-Fe7No,14711
-sleap_nn/train.py,sha256=XvVhzMXL9rNQLx1-6jIcp5BAO1pR7AZjdphMn5ZX-_I,27558
+sleap_nn/train.py,sha256=PEaK2B0S7DoImf8vt2cvJQS-n2NBw_pUJHmXy0J4NT0,30712
 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=l9lMJDxIMb-9MI3ShOtVwbOUMuwOLtSQlxiVyYHqjvE,13953
@@ -19,7 +19,7 @@ sleap_nn/config/__init__.py,sha256=l0xV1uJsGJfMPfWAqlUR7Ivu4cSCWsP-3Y9ueyPESuk,4
 sleap_nn/config/data_config.py,sha256=5a5YlXm4V9qGvkqgFNy6o0XJ_Q06UFjpYJXmNHfvXEI,24021
 sleap_nn/config/get_config.py,sha256=rjNUffKU9z-ohLwrOVmJNGCqwUM93eh68h4KJfrSy8Y,42396
 sleap_nn/config/model_config.py,sha256=XFIbqFno7IkX0Se5WF_2_7aUalAlC2SvpDe-uP2TttM,57582
-sleap_nn/config/trainer_config.py,sha256=ZMXxns6VYakgYHRhkM541Eje76DdaTdDi4FFPNjJtP4,28413
+sleap_nn/config/trainer_config.py,sha256=Ob2UqU10DXsQOnDb0iJxy0qc82CfP6FkQZQkrCvTEEY,29120
 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
@@ -34,11 +34,32 @@ sleap_nn/data/normalization.py,sha256=5xEvcguG-fvAGObl4nWPZ9TEM5gvv0uYPGDuni34XI
 sleap_nn/data/providers.py,sha256=0x6GFP1s1c08ji4p0M5V6p-dhT4Z9c-SI_Aw1DWX-uM,14272
 sleap_nn/data/resizing.py,sha256=YFpSQduIBkRK39FYmrqDL-v8zMySlEs6TJxh6zb_0ZU,5076
 sleap_nn/data/utils.py,sha256=rT0w7KMOTlzaeKWq1TqjbgC4Lvjz_G96McllvEOqXx8,5641
+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=NqN-G8TzAOsvCoL3bttEjA1iGsuveLOnOCXIUeFCdSA,13684
+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=L9LdYKt_Bfw7cxo6xEpgF8wXcZAwq5plCfmKJ839N40,13014
+sleap_nn/inference/postprocessing.py,sha256=ZM_OH7_WIprieaujZ2Rk_34JhSDDzCry6Pq2YM_u5sg,8998
 sleap_nn/inference/predictors.py,sha256=U114RlgOXKGm5iz1lnTfE3aN9S0WCh6gWhVP3KVewfc,158046
 sleap_nn/inference/provenance.py,sha256=0BekXyvpLMb0Vv6DjpctlLduG9RN-Q8jt5zDm783eZE,11204
 sleap_nn/inference/single_instance.py,sha256=rOns_5TsJ1rb-lwmHG3ZY-pOhXGN2D-SfW9RmBxxzcI,4089
@@ -52,14 +73,14 @@ sleap_nn/tracking/candidates/__init__.py,sha256=1O7NObIwshM7j1rLHmImbFphvkM9wY1j
 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=TVnQ6plNC2MnlTiY2rSCRuw2WRk5cQSziek_VPUcOEg,25994
-sleap_nn/training/lightning_modules.py,sha256=G3c4xJkYWW-iSRawzkgTqkGd4lTsbPiMTcB5Nvq7jes,85512
+sleap_nn/training/callbacks.py,sha256=7FvYLLoOxbjIu9v9zjDmzHZC4TW6y_P-N9J9GliHlSc,37944
+sleap_nn/training/lightning_modules.py,sha256=z98NBTrNy-GfCw4zatummJhVUO1fdjv_kPweAKcoaXc,108394
 sleap_nn/training/losses.py,sha256=gbdinUURh4QUzjmNd2UJpt4FXwecqKy9gHr65JZ1bZk,1632
-sleap_nn/training/model_trainer.py,sha256=loCmEX0DfBtdV_pN-W8s31fn2_L-lbpWaq3OQXeSp-0,59337
+sleap_nn/training/model_trainer.py,sha256=LHBIpjcI6WCJUxAHFKA-DSesNT3Nk4-RhSKlDaP4Fz4,58715
 sleap_nn/training/utils.py,sha256=ivdkZEI0DkTCm6NPszsaDOh9jSfozkONZdl6TvvQUWI,20398
-sleap_nn-0.1.0a1.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-sleap_nn-0.1.0a1.dist-info/METADATA,sha256=h3d4WPIu_JunY32jaRqJ4-fXp4KruTWT57FWb3L6dps,5637
-sleap_nn-0.1.0a1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sleap_nn-0.1.0a1.dist-info/entry_points.txt,sha256=zfl5Y3hidZxWBvo8qXvu5piJAXJ_l6v7xVFm0gNiUoI,46
-sleap_nn-0.1.0a1.dist-info/top_level.txt,sha256=Kz68iQ55K75LWgSeqz4V4SCMGeFFYH-KGBOyhQh3xZE,9
-sleap_nn-0.1.0a1.dist-info/RECORD,,
+sleap_nn-0.1.0a3.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+sleap_nn-0.1.0a3.dist-info/METADATA,sha256=nx3yU6SNyPf_hNqCO1183zj3OEuxX3frrMqIr0MfNbo,6145
+sleap_nn-0.1.0a3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sleap_nn-0.1.0a3.dist-info/entry_points.txt,sha256=zfl5Y3hidZxWBvo8qXvu5piJAXJ_l6v7xVFm0gNiUoI,46
+sleap_nn-0.1.0a3.dist-info/top_level.txt,sha256=Kz68iQ55K75LWgSeqz4V4SCMGeFFYH-KGBOyhQh3xZE,9
+sleap_nn-0.1.0a3.dist-info/RECORD,,