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

sleap_nn/training/lightning_modules.py

@@ -1,6 +1,6 @@
 """This module has the LightningModule classes for all model types."""
 
-from typing import Optional, Union, Dict, Any, List
+from typing import Optional, Union, Dict, Any
 import time
 from torch import nn
 import numpy as np
@@ -51,16 +51,10 @@ matplotlib.use(
 import matplotlib.pyplot as plt
 from sleap_nn.config.utils import get_backbone_type_from_cfg, get_model_type_from_cfg
 from sleap_nn.config.trainer_config import (
-    CosineAnnealingWarmupConfig,
-    LinearWarmupLinearDecayConfig,
     LRSchedulerConfig,
     ReduceLROnPlateauConfig,
     StepLRConfig,
 )
-from sleap_nn.training.schedulers import (
-    LinearWarmupCosineAnnealingLR,
-    LinearWarmupLinearDecayLR,
-)
 from sleap_nn.config.get_config import get_backbone_config
 from sleap_nn.legacy_models import (
     load_legacy_model_weights,
@@ -190,15 +184,6 @@ class LightningModel(L.LightningModule):
         self.val_loss = {}
         self.learning_rate = {}
 
-        # For epoch-averaged loss tracking
-        self._epoch_loss_sum = 0.0
-        self._epoch_loss_count = 0
-
-        # For epoch-end evaluation
-        self.val_predictions: List[Dict] = []
-        self.val_ground_truth: List[Dict] = []
-        self._collect_val_predictions: bool = False
-
         # Initialization for encoder and decoder stacks.
         if self.init_weights == "xavier":
             self.model.apply(xavier_init_weights)
@@ -235,9 +220,7 @@ class LightningModel(L.LightningModule):
             elif self.pretrained_backbone_weights.endswith(".h5"):
                 # load from sleap model weights
                 load_legacy_model_weights(
-                    self.model.backbone,
-                    self.pretrained_backbone_weights,
-                    component="backbone",
+                    self.model.backbone, self.pretrained_backbone_weights
                 )
 
             else:
@@ -266,9 +249,7 @@ class LightningModel(L.LightningModule):
             elif self.pretrained_head_weights.endswith(".h5"):
                 # load from sleap model weights
                 load_legacy_model_weights(
-                    self.model.head_layers,
-                    self.pretrained_head_weights,
-                    component="head",
+                    self.model.head_layers, self.pretrained_head_weights
                 )
 
             else:
@@ -324,24 +305,17 @@ class LightningModel(L.LightningModule):
     def on_train_epoch_start(self):
         """Configure the train timer at the beginning of each epoch."""
         self.train_start_time = time.time()
-        # Reset epoch loss tracking
-        self._epoch_loss_sum = 0.0
-        self._epoch_loss_count = 0
-
-    def _accumulate_loss(self, loss: torch.Tensor):
-        """Accumulate loss for epoch-averaged logging. Call this in training_step."""
-        self._epoch_loss_sum += loss.detach().item()
-        self._epoch_loss_count += 1
 
     def on_train_epoch_end(self):
         """Configure the train timer at the end of every epoch."""
         train_time = time.time() - self.train_start_time
         self.log(
-            "train/time",
+            "train_time",
             train_time,
             prog_bar=False,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         # Log epoch explicitly for custom x-axis support in wandb
@@ -350,56 +324,24 @@ class LightningModel(L.LightningModule):
             float(self.current_epoch),
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
-        # Log epoch-averaged training loss
-        if self._epoch_loss_count > 0:
-            avg_loss = self._epoch_loss_sum / self._epoch_loss_count
-            self.log(
-                "train/loss",
-                avg_loss,
-                prog_bar=False,
-                on_step=False,
-                on_epoch=True,
-                sync_dist=True,
-            )
-        # Log current learning rate (useful for monitoring LR schedulers)
-        if self.trainer.optimizers:
-            lr = self.trainer.optimizers[0].param_groups[0]["lr"]
-            self.log(
-                "train/lr",
-                lr,
-                prog_bar=False,
-                on_step=False,
-                on_epoch=True,
-                sync_dist=True,
-            )
 
     def on_validation_epoch_start(self):
         """Configure the val timer at the beginning of each epoch."""
         self.val_start_time = time.time()
-        # Clear accumulated predictions for new epoch
-        self.val_predictions = []
-        self.val_ground_truth = []
 
     def on_validation_epoch_end(self):
         """Configure the val timer at the end of every epoch."""
         val_time = time.time() - self.val_start_time
         self.log(
-            "val/time",
+            "val_time",
             val_time,
             prog_bar=False,
             on_step=False,
             on_epoch=True,
-            sync_dist=True,
-        )
-        # Log epoch explicitly so val/* metrics can use it as x-axis in wandb
-        # (mirrors what on_train_epoch_end does for train/* metrics)
-        self.log(
-            "epoch",
-            float(self.current_epoch),
-            on_step=False,
-            on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
 
@@ -436,51 +378,13 @@ class LightningModel(L.LightningModule):
                 lr_scheduler_cfg.step_lr = StepLRConfig()
             elif self.lr_scheduler == "reduce_lr_on_plateau":
                 lr_scheduler_cfg.reduce_lr_on_plateau = ReduceLROnPlateauConfig()
-            elif self.lr_scheduler == "cosine_annealing_warmup":
-                lr_scheduler_cfg.cosine_annealing_warmup = CosineAnnealingWarmupConfig()
-            elif self.lr_scheduler == "linear_warmup_linear_decay":
-                lr_scheduler_cfg.linear_warmup_linear_decay = (
-                    LinearWarmupLinearDecayConfig()
-                )
 
         elif isinstance(self.lr_scheduler, dict):
             lr_scheduler_cfg = self.lr_scheduler
 
         for k, v in self.lr_scheduler.items():
             if v is not None:
-                if k == "cosine_annealing_warmup":
-                    cfg = self.lr_scheduler.cosine_annealing_warmup
-                    # Use trainer's max_epochs if not specified in config
-                    max_epochs = (
-                        cfg.max_epochs
-                        if cfg.max_epochs is not None
-                        else self.trainer.max_epochs
-                    )
-                    scheduler = LinearWarmupCosineAnnealingLR(
-                        optimizer=optimizer,
-                        warmup_epochs=cfg.warmup_epochs,
-                        max_epochs=max_epochs,
-                        warmup_start_lr=cfg.warmup_start_lr,
-                        eta_min=cfg.eta_min,
-                    )
-                    break
-                elif k == "linear_warmup_linear_decay":
-                    cfg = self.lr_scheduler.linear_warmup_linear_decay
-                    # Use trainer's max_epochs if not specified in config
-                    max_epochs = (
-                        cfg.max_epochs
-                        if cfg.max_epochs is not None
-                        else self.trainer.max_epochs
-                    )
-                    scheduler = LinearWarmupLinearDecayLR(
-                        optimizer=optimizer,
-                        warmup_epochs=cfg.warmup_epochs,
-                        max_epochs=max_epochs,
-                        warmup_start_lr=cfg.warmup_start_lr,
-                        end_lr=cfg.end_lr,
-                    )
-                    break
-                elif k == "step_lr":
+                if k == "step_lr":
                     scheduler = torch.optim.lr_scheduler.StepLR(
                         optimizer=optimizer,
                         step_size=self.lr_scheduler.step_lr.step_size,
@@ -508,7 +412,7 @@ class LightningModel(L.LightningModule):
             "optimizer": optimizer,
             "lr_scheduler": {
                 "scheduler": scheduler,
-                "monitor": "val/loss",
+                "monitor": "val_loss",
             },
         }
 
@@ -664,7 +568,6 @@ class SingleInstanceLightningModule(LightningModel):
             torch.squeeze(batch["image"], dim=1),
             torch.squeeze(batch["confidence_maps"], dim=1),
         )
-        X = normalize_on_gpu(X)
 
         y_preds = self.model(X)["SingleInstanceConfmapsHead"]
 
@@ -688,24 +591,23 @@ class SingleInstanceLightningModule(LightningModel):
             channel_wise_loss = torch.sum(mse, dim=(0, 2, 3)) / (batch_size * h * w)
             for node_idx, name in enumerate(self.node_names):
                 self.log(
-                    f"train/confmaps/{name}",
+                    f"{name}",
                     channel_wise_loss[node_idx],
                     prog_bar=False,
                     on_step=False,
                     on_epoch=True,
+                    logger=True,
                     sync_dist=True,
                 )
-        # Log step-level loss (every batch, uses global_step x-axis)
         self.log(
-            "loss",
+            "train_loss",
             loss,
             prog_bar=True,
             on_step=True,
             on_epoch=False,
+            logger=True,
             sync_dist=True,
         )
-        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
-        self._accumulate_loss(loss)
         return loss
 
     def validation_step(self, batch, batch_idx):
@@ -714,7 +616,6 @@ class SingleInstanceLightningModule(LightningModel):
             torch.squeeze(batch["image"], dim=1),
             torch.squeeze(batch["confidence_maps"], dim=1),
         )
-        X = normalize_on_gpu(X)
 
         y_preds = self.model(X)["SingleInstanceConfmapsHead"]
         val_loss = nn.MSELoss()(y_preds, y)
@@ -729,59 +630,15 @@ class SingleInstanceLightningModule(LightningModel):
             )
             val_loss = val_loss + ohkm_loss
         self.log(
-            "val/loss",
+            "val_loss",
             val_loss,
             prog_bar=True,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
 
-        # Collect predictions for epoch-end evaluation if enabled
-        if self._collect_val_predictions:
-            with torch.no_grad():
-                # Squeeze n_samples dim from image for inference (batch, 1, C, H, W) -> (batch, C, H, W)
-                inference_batch = {k: v for k, v in batch.items()}
-                if inference_batch["image"].ndim == 5:
-                    inference_batch["image"] = inference_batch["image"].squeeze(1)
-                inference_output = self.single_instance_inf_layer(inference_batch)
-                if isinstance(inference_output, list):
-                    inference_output = inference_output[0]
-
-            batch_size = len(batch["frame_idx"])
-            for i in range(batch_size):
-                eff = batch["eff_scale"][i].cpu().numpy()
-
-                # Predictions are already in original image space (inference divides by eff_scale)
-                pred_peaks = inference_output["pred_instance_peaks"][i].cpu().numpy()
-                pred_scores = inference_output["pred_peak_values"][i].cpu().numpy()
-
-                # Transform GT from preprocessed to original image space
-                # Note: instances have shape (1, max_inst, n_nodes, 2) - squeeze n_samples dim
-                gt_prep = batch["instances"][i].cpu().numpy()
-                if gt_prep.ndim == 4:
-                    gt_prep = gt_prep.squeeze(0)  # (max_inst, n_nodes, 2)
-                gt_orig = gt_prep / eff
-                num_inst = batch["num_instances"][i].item()
-                gt_orig = gt_orig[:num_inst]  # Only valid instances
-
-                self.val_predictions.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "pred_peaks": pred_peaks,
-                        "pred_scores": pred_scores,
-                    }
-                )
-                self.val_ground_truth.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "gt_instances": gt_orig,
-                        "num_instances": num_inst,
-                    }
-                )
-
 
 class TopDownCenteredInstanceLightningModule(LightningModel):
     """Lightning Module for TopDownCenteredInstance Model.
@@ -927,7 +784,6 @@ class TopDownCenteredInstanceLightningModule(LightningModel):
             torch.squeeze(batch["instance_image"], dim=1),
             torch.squeeze(batch["confidence_maps"], dim=1),
         )
-        X = normalize_on_gpu(X)
 
         y_preds = self.model(X)["CenteredInstanceConfmapsHead"]
 
@@ -951,25 +807,24 @@ class TopDownCenteredInstanceLightningModule(LightningModel):
             channel_wise_loss = torch.sum(mse, dim=(0, 2, 3)) / (batch_size * h * w)
             for node_idx, name in enumerate(self.node_names):
                 self.log(
-                    f"train/confmaps/{name}",
+                    f"{name}",
                     channel_wise_loss[node_idx],
                     prog_bar=False,
                     on_step=False,
                     on_epoch=True,
+                    logger=True,
                     sync_dist=True,
                 )
 
-        # Log step-level loss (every batch, uses global_step x-axis)
         self.log(
-            "loss",
+            "train_loss",
             loss,
             prog_bar=True,
             on_step=True,
             on_epoch=False,
+            logger=True,
             sync_dist=True,
         )
-        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
-        self._accumulate_loss(loss)
         return loss
 
     def validation_step(self, batch, batch_idx):
@@ -978,7 +833,6 @@ class TopDownCenteredInstanceLightningModule(LightningModel):
             torch.squeeze(batch["instance_image"], dim=1),
             torch.squeeze(batch["confidence_maps"], dim=1),
         )
-        X = normalize_on_gpu(X)
 
         y_preds = self.model(X)["CenteredInstanceConfmapsHead"]
         val_loss = nn.MSELoss()(y_preds, y)
@@ -993,70 +847,15 @@ class TopDownCenteredInstanceLightningModule(LightningModel):
             )
             val_loss = val_loss + ohkm_loss
         self.log(
-            "val/loss",
+            "val_loss",
             val_loss,
             prog_bar=True,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
 
-        # Collect predictions for epoch-end evaluation if enabled
-        if self._collect_val_predictions:
-            # SAVE bbox BEFORE inference (it modifies in-place!)
-            bbox_prep_saved = batch["instance_bbox"].clone()
-
-            with torch.no_grad():
-                inference_output = self.instance_peaks_inf_layer(batch)
-
-            batch_size = len(batch["frame_idx"])
-            for i in range(batch_size):
-                eff = batch["eff_scale"][i].cpu().numpy()
-
-                # Predictions from inference (crop-relative, original scale)
-                pred_peaks_crop = (
-                    inference_output["pred_instance_peaks"][i].cpu().numpy()
-                )
-                pred_scores = inference_output["pred_peak_values"][i].cpu().numpy()
-
-                # Compute bbox offset in original space from SAVED prep bbox
-                # bbox has shape (n_samples=1, 4, 2) where 4 corners
-                bbox_prep = bbox_prep_saved[i].squeeze(0).cpu().numpy()  # (4, 2)
-                bbox_top_left_orig = (
-                    bbox_prep[0] / eff
-                )  # Top-left corner in original space
-
-                # Full image coordinates (original space)
-                pred_peaks_full = pred_peaks_crop + bbox_top_left_orig
-
-                # GT transform: crop-relative preprocessed -> full image original
-                gt_crop_prep = (
-                    batch["instance"][i].squeeze(0).cpu().numpy()
-                )  # (n_nodes, 2)
-                gt_crop_orig = gt_crop_prep / eff
-                gt_full_orig = gt_crop_orig + bbox_top_left_orig
-
-                self.val_predictions.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "pred_peaks": pred_peaks_full.reshape(
-                            1, -1, 2
-                        ),  # (1, n_nodes, 2)
-                        "pred_scores": pred_scores.reshape(1, -1),  # (1, n_nodes)
-                    }
-                )
-                self.val_ground_truth.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "gt_instances": gt_full_orig.reshape(
-                            1, -1, 2
-                        ),  # (1, n_nodes, 2)
-                        "num_instances": 1,
-                    }
-                )
-
 
 class CentroidLightningModule(LightningModel):
     """Lightning Module for Centroid Model.
@@ -1202,21 +1001,18 @@ class CentroidLightningModule(LightningModel):
             torch.squeeze(batch["image"], dim=1),
             torch.squeeze(batch["centroids_confidence_maps"], dim=1),
         )
-        X = normalize_on_gpu(X)
 
         y_preds = self.model(X)["CentroidConfmapsHead"]
         loss = nn.MSELoss()(y_preds, y)
-        # Log step-level loss (every batch, uses global_step x-axis)
         self.log(
-            "loss",
+            "train_loss",
             loss,
             prog_bar=True,
             on_step=True,
             on_epoch=False,
+            logger=True,
             sync_dist=True,
         )
-        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
-        self._accumulate_loss(loss)
         return loss
 
     def validation_step(self, batch, batch_idx):
@@ -1225,74 +1021,19 @@ class CentroidLightningModule(LightningModel):
             torch.squeeze(batch["image"], dim=1),
             torch.squeeze(batch["centroids_confidence_maps"], dim=1),
         )
-        X = normalize_on_gpu(X)
 
         y_preds = self.model(X)["CentroidConfmapsHead"]
         val_loss = nn.MSELoss()(y_preds, y)
         self.log(
-            "val/loss",
+            "val_loss",
             val_loss,
             prog_bar=True,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
 
-        # Collect predictions for epoch-end evaluation if enabled
-        if self._collect_val_predictions:
-            # Save GT centroids before inference (inference overwrites batch["centroids"])
-            batch["gt_centroids"] = batch["centroids"].clone()
-
-            with torch.no_grad():
-                inference_output = self.centroid_inf_layer(batch)
-
-            batch_size = len(batch["frame_idx"])
-            for i in range(batch_size):
-                eff = batch["eff_scale"][i].cpu().numpy()
-
-                # Predictions are in original image space (inference divides by eff_scale)
-                # centroids shape: (batch, 1, max_instances, 2) - squeeze to (max_instances, 2)
-                pred_centroids = (
-                    inference_output["centroids"][i].squeeze(0).cpu().numpy()
-                )
-                pred_vals = inference_output["centroid_vals"][i].cpu().numpy()
-
-                # Transform GT centroids from preprocessed to original image space
-                # Use "gt_centroids" since inference overwrites "centroids" with predictions
-                gt_centroids_prep = (
-                    batch["gt_centroids"][i].cpu().numpy()
-                )  # (n_samples=1, max_inst, 2)
-                gt_centroids_orig = gt_centroids_prep.squeeze(0) / eff  # (max_inst, 2)
-                num_inst = batch["num_instances"][i].item()
-
-                # Filter to valid instances (non-NaN)
-                valid_pred_mask = ~np.isnan(pred_centroids).any(axis=1)
-                pred_centroids = pred_centroids[valid_pred_mask]
-                pred_vals = pred_vals[valid_pred_mask]
-
-                gt_centroids_valid = gt_centroids_orig[:num_inst]
-
-                self.val_predictions.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "pred_peaks": pred_centroids.reshape(
-                            -1, 1, 2
-                        ),  # (n_inst, 1, 2)
-                        "pred_scores": pred_vals.reshape(-1, 1),  # (n_inst, 1)
-                    }
-                )
-                self.val_ground_truth.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "gt_instances": gt_centroids_valid.reshape(
-                            -1, 1, 2
-                        ),  # (n_inst, 1, 2)
-                        "num_instances": num_inst,
-                    }
-                )
-
 
 class BottomUpLightningModule(LightningModel):
     """Lightning Module for BottomUp Model.
@@ -1385,13 +1126,12 @@ class BottomUpLightningModule(LightningModel):
         self.bottomup_inf_layer = BottomUpInferenceModel(
             torch_model=self.forward,
             paf_scorer=paf_scorer,
-            peak_threshold=0.1,  # Lower threshold for epoch-end eval during training
+            peak_threshold=0.2,
             input_scale=1.0,
             return_confmaps=True,
             return_pafs=True,
             cms_output_stride=self.head_configs.bottomup.confmaps.output_stride,
             pafs_output_stride=self.head_configs.bottomup.pafs.output_stride,
-            max_peaks_per_node=100,  # Prevents combinatorial explosion in early training
         )
         self.node_names = list(self.head_configs.bottomup.confmaps.part_names)
 
@@ -1476,7 +1216,6 @@ class BottomUpLightningModule(LightningModel):
         X = torch.squeeze(batch["image"], dim=1)
         y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
         y_paf = batch["part_affinity_fields"]
-        X = normalize_on_gpu(X)
         preds = self.model(X)
         pafs = preds["PartAffinityFieldsHead"]
         confmaps = preds["MultiInstanceConfmapsHead"]
@@ -1509,29 +1248,29 @@ class BottomUpLightningModule(LightningModel):
             "PartAffinityFieldsHead": pafs_loss,
         }
         loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
-        # Log step-level loss (every batch, uses global_step x-axis)
         self.log(
-            "loss",
+            "train_loss",
             loss,
             prog_bar=True,
             on_step=True,
             on_epoch=False,
+            logger=True,
             sync_dist=True,
         )
-        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
-        self._accumulate_loss(loss)
         self.log(
-            "train/confmaps_loss",
+            "train_confmap_loss",
             confmap_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         self.log(
-            "train/paf_loss",
+            "train_paf_loss",
             pafs_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         return loss
@@ -1541,7 +1280,6 @@ class BottomUpLightningModule(LightningModel):
         X = torch.squeeze(batch["image"], dim=1)
         y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
         y_paf = batch["part_affinity_fields"]
-        X = normalize_on_gpu(X)
 
         preds = self.model(X)
         pafs = preds["PartAffinityFieldsHead"]
@@ -1577,75 +1315,31 @@ class BottomUpLightningModule(LightningModel):
 
         val_loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
         self.log(
-            "val/loss",
+            "val_loss",
             val_loss,
             prog_bar=True,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         self.log(
-            "val/confmaps_loss",
+            "val_confmap_loss",
             confmap_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         self.log(
-            "val/paf_loss",
+            "val_paf_loss",
             pafs_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
 
-        # Collect predictions for epoch-end evaluation if enabled
-        if self._collect_val_predictions:
-            with torch.no_grad():
-                # Note: Do NOT squeeze the image here - the forward() method expects
-                # (batch, n_samples, C, H, W) and handles the n_samples squeeze internally
-                inference_output = self.bottomup_inf_layer(batch)
-                if isinstance(inference_output, list):
-                    inference_output = inference_output[0]
-
-            batch_size = len(batch["frame_idx"])
-            for i in range(batch_size):
-                eff = batch["eff_scale"][i].cpu().numpy()
-
-                # Predictions are already in original space (variable number of instances)
-                pred_peaks = inference_output["pred_instance_peaks"][i]
-                pred_scores = inference_output["pred_peak_values"][i]
-                if torch.is_tensor(pred_peaks):
-                    pred_peaks = pred_peaks.cpu().numpy()
-                if torch.is_tensor(pred_scores):
-                    pred_scores = pred_scores.cpu().numpy()
-
-                # Transform GT to original space
-                # Note: instances have shape (1, max_inst, n_nodes, 2) - squeeze n_samples dim
-                gt_prep = batch["instances"][i].cpu().numpy()
-                if gt_prep.ndim == 4:
-                    gt_prep = gt_prep.squeeze(0)  # (max_inst, n_nodes, 2)
-                gt_orig = gt_prep / eff
-                num_inst = batch["num_instances"][i].item()
-                gt_orig = gt_orig[:num_inst]  # Only valid instances
-
-                self.val_predictions.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "pred_peaks": pred_peaks,  # Original space, variable instances
-                        "pred_scores": pred_scores,
-                    }
-                )
-                self.val_ground_truth.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "gt_instances": gt_orig,  # Original space
-                        "num_instances": num_inst,
-                    }
-                )
-
 
 class BottomUpMultiClassLightningModule(LightningModel):
     """Lightning Module for BottomUp ID Model.
@@ -1824,7 +1518,6 @@ class BottomUpMultiClassLightningModule(LightningModel):
         X = torch.squeeze(batch["image"], dim=1)
         y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
         y_classmap = torch.squeeze(batch["class_maps"], dim=1)
-        X = normalize_on_gpu(X)
         preds = self.model(X)
         classmaps = preds["ClassMapsHead"]
         confmaps = preds["MultiInstanceConfmapsHead"]
@@ -1848,84 +1541,31 @@ class BottomUpMultiClassLightningModule(LightningModel):
             "ClassMapsHead": classmaps_loss,
         }
         loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
-        # Log step-level loss (every batch, uses global_step x-axis)
         self.log(
-            "loss",
+            "train_loss",
             loss,
             prog_bar=True,
             on_step=True,
             on_epoch=False,
+            logger=True,
             sync_dist=True,
         )
-        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
-        self._accumulate_loss(loss)
         self.log(
-            "train/confmaps_loss",
+            "train_confmap_loss",
             confmap_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         self.log(
-            "train/classmap_loss",
+            "train_classmap_loss",
             classmaps_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
-
-        # Compute classification accuracy at GT keypoint locations
-        with torch.no_grad():
-            # Get output stride for class maps
-            cms_stride = self.head_configs.multi_class_bottomup.class_maps.output_stride
-
-            # Get GT instances and sample class maps at those locations
-            instances = batch["instances"]  # (batch, n_samples, max_inst, n_nodes, 2)
-            if instances.dim() == 5:
-                instances = instances.squeeze(1)  # (batch, max_inst, n_nodes, 2)
-            num_instances = batch["num_instances"]  # (batch,)
-
-            correct = 0
-            total = 0
-            for b in range(instances.shape[0]):
-                n_inst = num_instances[b].item()
-                for inst_idx in range(n_inst):
-                    for node_idx in range(instances.shape[2]):
-                        # Get keypoint location (in input image space)
-                        kp = instances[b, inst_idx, node_idx]  # (2,) = (x, y)
-                        if torch.isnan(kp).any():
-                            continue
-
-                        # Convert to class map space
-                        x_cm = (
-                            (kp[0] / cms_stride)
-                            .long()
-                            .clamp(0, classmaps.shape[-1] - 1)
-                        )
-                        y_cm = (
-                            (kp[1] / cms_stride)
-                            .long()
-                            .clamp(0, classmaps.shape[-2] - 1)
-                        )
-
-                        # Sample predicted and GT class at this location
-                        pred_class = classmaps[b, :, y_cm, x_cm].argmax()
-                        gt_class = y_classmap[b, :, y_cm, x_cm].argmax()
-
-                        if pred_class == gt_class:
-                            correct += 1
-                        total += 1
-
-            if total > 0:
-                class_accuracy = torch.tensor(correct / total, device=X.device)
-                self.log(
-                    "train/class_accuracy",
-                    class_accuracy,
-                    on_step=False,
-                    on_epoch=True,
-                    sync_dist=True,
-                )
-
         return loss
 
     def validation_step(self, batch, batch_idx):
@@ -1933,7 +1573,6 @@ class BottomUpMultiClassLightningModule(LightningModel):
         X = torch.squeeze(batch["image"], dim=1)
         y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
         y_classmap = torch.squeeze(batch["class_maps"], dim=1)
-        X = normalize_on_gpu(X)
 
         preds = self.model(X)
         classmaps = preds["ClassMapsHead"]
@@ -1960,127 +1599,31 @@ class BottomUpMultiClassLightningModule(LightningModel):
 
         val_loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
         self.log(
-            "val/loss",
+            "val_loss",
             val_loss,
             prog_bar=True,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         self.log(
-            "val/confmaps_loss",
+            "val_confmap_loss",
             confmap_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         self.log(
-            "val/classmap_loss",
+            "val_classmap_loss",
             classmaps_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
 
-        # Compute classification accuracy at GT keypoint locations
-        with torch.no_grad():
-            # Get output stride for class maps
-            cms_stride = self.head_configs.multi_class_bottomup.class_maps.output_stride
-
-            # Get GT instances and sample class maps at those locations
-            instances = batch["instances"]  # (batch, n_samples, max_inst, n_nodes, 2)
-            if instances.dim() == 5:
-                instances = instances.squeeze(1)  # (batch, max_inst, n_nodes, 2)
-            num_instances = batch["num_instances"]  # (batch,)
-
-            correct = 0
-            total = 0
-            for b in range(instances.shape[0]):
-                n_inst = num_instances[b].item()
-                for inst_idx in range(n_inst):
-                    for node_idx in range(instances.shape[2]):
-                        # Get keypoint location (in input image space)
-                        kp = instances[b, inst_idx, node_idx]  # (2,) = (x, y)
-                        if torch.isnan(kp).any():
-                            continue
-
-                        # Convert to class map space
-                        x_cm = (
-                            (kp[0] / cms_stride)
-                            .long()
-                            .clamp(0, classmaps.shape[-1] - 1)
-                        )
-                        y_cm = (
-                            (kp[1] / cms_stride)
-                            .long()
-                            .clamp(0, classmaps.shape[-2] - 1)
-                        )
-
-                        # Sample predicted and GT class at this location
-                        pred_class = classmaps[b, :, y_cm, x_cm].argmax()
-                        gt_class = y_classmap[b, :, y_cm, x_cm].argmax()
-
-                        if pred_class == gt_class:
-                            correct += 1
-                        total += 1
-
-            if total > 0:
-                class_accuracy = torch.tensor(correct / total, device=X.device)
-                self.log(
-                    "val/class_accuracy",
-                    class_accuracy,
-                    on_step=False,
-                    on_epoch=True,
-                    sync_dist=True,
-                )
-
-        # Collect predictions for epoch-end evaluation if enabled
-        if self._collect_val_predictions:
-            with torch.no_grad():
-                # Note: Do NOT squeeze the image here - the forward() method expects
-                # (batch, n_samples, C, H, W) and handles the n_samples squeeze internally
-                inference_output = self.bottomup_inf_layer(batch)
-                if isinstance(inference_output, list):
-                    inference_output = inference_output[0]
-
-            batch_size = len(batch["frame_idx"])
-            for i in range(batch_size):
-                eff = batch["eff_scale"][i].cpu().numpy()
-
-                # Predictions are already in original space (variable number of instances)
-                pred_peaks = inference_output["pred_instance_peaks"][i]
-                pred_scores = inference_output["pred_peak_values"][i]
-                if torch.is_tensor(pred_peaks):
-                    pred_peaks = pred_peaks.cpu().numpy()
-                if torch.is_tensor(pred_scores):
-                    pred_scores = pred_scores.cpu().numpy()
-
-                # Transform GT to original space
-                # Note: instances have shape (1, max_inst, n_nodes, 2) - squeeze n_samples dim
-                gt_prep = batch["instances"][i].cpu().numpy()
-                if gt_prep.ndim == 4:
-                    gt_prep = gt_prep.squeeze(0)  # (max_inst, n_nodes, 2)
-                gt_orig = gt_prep / eff
-                num_inst = batch["num_instances"][i].item()
-                gt_orig = gt_orig[:num_inst]  # Only valid instances
-
-                self.val_predictions.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "pred_peaks": pred_peaks,  # Original space, variable instances
-                        "pred_scores": pred_scores,
-                    }
-                )
-                self.val_ground_truth.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "gt_instances": gt_orig,  # Original space
-                        "num_instances": num_inst,
-                    }
-                )
-
 
 class TopDownCenteredInstanceMultiClassLightningModule(LightningModel):
     """Lightning Module for TopDownCenteredInstance ID Model.
@@ -2229,7 +1772,6 @@ class TopDownCenteredInstanceMultiClassLightningModule(LightningModel):
         X = torch.squeeze(batch["instance_image"], dim=1)
         y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
         y_classvector = batch["class_vectors"]
-        X = normalize_on_gpu(X)
         preds = self.model(X)
         classvector = preds["ClassVectorsHead"]
         confmaps = preds["CenteredInstanceConfmapsHead"]
@@ -2261,50 +1803,38 @@ class TopDownCenteredInstanceMultiClassLightningModule(LightningModel):
             channel_wise_loss = torch.sum(mse, dim=(0, 2, 3)) / (batch_size * h * w)
             for node_idx, name in enumerate(self.node_names):
                 self.log(
-                    f"train/confmaps/{name}",
+                    f"{name}",
                     channel_wise_loss[node_idx],
                     prog_bar=False,
                     on_step=False,
                     on_epoch=True,
+                    logger=True,
                     sync_dist=True,
                 )
 
-        # Log step-level loss (every batch, uses global_step x-axis)
         self.log(
-            "loss",
+            "train_loss",
             loss,
             prog_bar=True,
             on_step=True,
             on_epoch=False,
+            logger=True,
             sync_dist=True,
         )
-        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
-        self._accumulate_loss(loss)
         self.log(
-            "train/confmaps_loss",
+            "train_confmap_loss",
             confmap_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         self.log(
-            "train/classvector_loss",
+            "train_classvector_loss",
             classvector_loss,
             on_step=False,
             on_epoch=True,
-            sync_dist=True,
-        )
-
-        # Compute classification accuracy
-        with torch.no_grad():
-            pred_classes = torch.argmax(classvector, dim=1)
-            gt_classes = torch.argmax(y_classvector, dim=1)
-            class_accuracy = (pred_classes == gt_classes).float().mean()
-        self.log(
-            "train/class_accuracy",
-            class_accuracy,
-            on_step=False,
-            on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         return loss
@@ -2314,7 +1844,6 @@ class TopDownCenteredInstanceMultiClassLightningModule(LightningModel):
         X = torch.squeeze(batch["instance_image"], dim=1)
         y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
         y_classvector = batch["class_vectors"]
-        X = normalize_on_gpu(X)
         preds = self.model(X)
         classvector = preds["ClassVectorsHead"]
         confmaps = preds["CenteredInstanceConfmapsHead"]
@@ -2339,93 +1868,27 @@ class TopDownCenteredInstanceMultiClassLightningModule(LightningModel):
         }
         val_loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
         self.log(
-            "val/loss",
+            "val_loss",
             val_loss,
             prog_bar=True,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         self.log(
-            "val/confmaps_loss",
+            "val_confmap_loss",
             confmap_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
         self.log(
-            "val/classvector_loss",
+            "val_classvector_loss",
             classvector_loss,
             on_step=False,
             on_epoch=True,
+            logger=True,
             sync_dist=True,
         )
-
-        # Compute classification accuracy
-        with torch.no_grad():
-            pred_classes = torch.argmax(classvector, dim=1)
-            gt_classes = torch.argmax(y_classvector, dim=1)
-            class_accuracy = (pred_classes == gt_classes).float().mean()
-        self.log(
-            "val/class_accuracy",
-            class_accuracy,
-            on_step=False,
-            on_epoch=True,
-            sync_dist=True,
-        )
-
-        # Collect predictions for epoch-end evaluation if enabled
-        if self._collect_val_predictions:
-            # SAVE bbox BEFORE inference (it modifies in-place!)
-            bbox_prep_saved = batch["instance_bbox"].clone()
-
-            with torch.no_grad():
-                inference_output = self.instance_peaks_inf_layer(batch)
-
-            batch_size = len(batch["frame_idx"])
-            for i in range(batch_size):
-                eff = batch["eff_scale"][i].cpu().numpy()
-
-                # Predictions from inference (crop-relative, original scale)
-                pred_peaks_crop = (
-                    inference_output["pred_instance_peaks"][i].cpu().numpy()
-                )
-                pred_scores = inference_output["pred_peak_values"][i].cpu().numpy()
-
-                # Compute bbox offset in original space from SAVED prep bbox
-                # bbox has shape (n_samples=1, 4, 2) where 4 corners
-                bbox_prep = bbox_prep_saved[i].squeeze(0).cpu().numpy()  # (4, 2)
-                bbox_top_left_orig = (
-                    bbox_prep[0] / eff
-                )  # Top-left corner in original space
-
-                # Full image coordinates (original space)
-                pred_peaks_full = pred_peaks_crop + bbox_top_left_orig
-
-                # GT transform: crop-relative preprocessed -> full image original
-                gt_crop_prep = (
-                    batch["instance"][i].squeeze(0).cpu().numpy()
-                )  # (n_nodes, 2)
-                gt_crop_orig = gt_crop_prep / eff
-                gt_full_orig = gt_crop_orig + bbox_top_left_orig
-
-                self.val_predictions.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "pred_peaks": pred_peaks_full.reshape(
-                            1, -1, 2
-                        ),  # (1, n_nodes, 2)
-                        "pred_scores": pred_scores.reshape(1, -1),  # (1, n_nodes)
-                    }
-                )
-                self.val_ground_truth.append(
-                    {
-                        "video_idx": batch["video_idx"][i].item(),
-                        "frame_idx": batch["frame_idx"][i].item(),
-                        "gt_instances": gt_full_orig.reshape(
-                            1, -1, 2
-                        ),  # (1, n_nodes, 2)
-                        "num_instances": 1,
-                    }
-                )