From 678c780907f116dfa7e8e5eaaa5fd429429713fe Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Wed, 22 Jul 2026 20:19:51 +0000
Subject: [PATCH 01/10] refactor: val every k steps

---
 .../machine_learning/data_handling.py         |   2 +
 .../machine_learning/train.py                 | 162 +++++++++---------
 2 files changed, 79 insertions(+), 85 deletions(-)

diff --git a/src/aind_exaspim_image_compression/machine_learning/data_handling.py b/src/aind_exaspim_image_compression/machine_learning/data_handling.py
index 6d430ed..242131b 100644
--- a/src/aind_exaspim_image_compression/machine_learning/data_handling.py
+++ b/src/aind_exaspim_image_compression/machine_learning/data_handling.py
@@ -77,6 +77,8 @@ def build_training_example(
         transform.forward(target),
         fg.astype(np.float32),
     )
+
+
 from aind_exaspim_image_compression.utils.swc_util import Reader
 
 
diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index d3dce49..6d5adfd 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -16,7 +16,6 @@
 import numpy as np
 import os
 import torch
-import torch.nn as nn
 import torch.optim as optim
 from skimage import io
 
@@ -75,7 +74,7 @@ def __init__(
         fg_weight=20.0,
         num_workers=None,
         prefetch=2,
-        val_every=1,
+        val_every=1000,
         seed=0,
     ):
         """
@@ -103,10 +102,11 @@ def __init__(
             checkpoint selection. Default is False so validation matches FP32
             production inference.
         val_every : int, optional
-            Run validation (and checkpoint selection) every this many epochs;
-            the final epoch is always validated. The count-space metrics are
-            CPU-bound, so a large validation set is only cheap if it is not run
-            every epoch. Default is 1 (validate every epoch).
+            Run validation (and checkpoint selection) every this many gradient
+            updates; the final epoch is always validated. Count-space metrics
+            are CPU-bound, so a large validation set is only cheap if it is
+            not run every epoch. Default is 1000 (validate every 1000 gradient
+            updates).
         seed : int, optional
             Seed used to reproducibly shuffle training examples by epoch.
             Default is 0.
@@ -128,17 +128,16 @@ def __init__(
         self.use_amp = bool(use_amp)
         self.use_amp_validation = bool(use_amp_validation)
 
-        self.codec = blosc.Blosc(cname="zstd", clevel=5, shuffle=blosc.SHUFFLE)
+        self.codec = blosc.Blosc(cname="zstd", clevel=6, shuffle=blosc.SHUFFLE)
         self.criterion = SignalPreservingLoss(fg_weight=fg_weight)
         self.checkpoint_weights = checkpoint_weights
-        self.best_score = np.inf
         self.model = model.to(device) if model else UNet().to(device)
         self._resume_transform_cfg = None
         self.optimizer = optim.AdamW(self.model.parameters(), lr=lr)
         # T_max spans the whole run so the cosine anneals once. With a small
         # T_max the LR returns to its peak every 2*T_max epochs, and each
         # return destabilized training (growing periodic loss spikes).
-        self.scheduler = CosineAnnealingLR(self.optimizer, T_max=max_epochs)
+        self.scheduler = CosineAnnealingLR(self.optimizer, T_max=self.val_every)
         self.writer = SummaryWriter(log_dir=log_dir)
 
         # Scale the loss before backward so small float16 gradients do not
@@ -173,6 +172,7 @@ def run(self, train_dataset, val_dataset):
                     "resume checkpoint transform does not match the validation "
                     "dataset transform"
                 )
+
         self.transform = train_dataset.transform
         train_dataloader = DataLoader(
             train_dataset,
@@ -191,74 +191,75 @@ def run(self, train_dataset, val_dataset):
         )
 
         # Main
-        self.best_score = np.inf
+        step = 0
+        best_score = np.inf
+        self.model.train()
         for epoch in range(self.max_epochs):
-            # Train
+            # Epoch initializations
+            train_losses = []
             train_dataloader.set_epoch(epoch)
-            train_loss = self.train_step(train_dataloader, epoch)
-
-            # Validate every val_every epochs (and always on the final epoch);
-            # the count-space metrics are CPU-bound, so a large validation set
-            # is only cheap when it is not run every epoch.
-            is_last = epoch == self.max_epochs - 1
-            if epoch % self.val_every == 0 or is_last:
-                val_loss, val_cratio, is_best = self.validate_step(
-                    val_dataloader, epoch
-                )
-                suffix = " - New Best!" if is_best else ""
-                s = f"Epoch {epoch}:  train_loss={train_loss},  val_loss={val_loss}, val_cratio={val_cratio}" + suffix
-            else:
-                s = f"Epoch {epoch}:  train_loss={train_loss}"
-            print(s)
 
-            # Step scheduler
-            self.scheduler.step()
-
-    def train_step(self, train_dataloader, epoch):
+            # Run train/val epoch
+            for x, y, fg_mask in train_dataloader:
+                # Train
+                train_loss = self.train_step(x, y, fg_mask)
+                train_losses.append(train_loss)
+                self.scheduler.step()
+                step += 1
+
+                # Validate (if applicable)
+                if step % self.val_every == 0:
+                    val_loss, val_cratio, val_score = self.validate(
+                        val_dataloader,
+                        step,
+                    )
+                    is_best = epoch > 0 and val_score < best_score
+                    if is_best:
+                        best_score = val_score
+                    if epoch > 0:
+                        self.save_model(step, val_score)
+
+                # Report results (if applicable)
+                if step % self.val_every == 0:
+                    avg_train_loss = np.mean(train_losses)
+                    self.writer.add_scalar(
+                        "train_loss", avg_train_loss, step
+                    )
+
+                    suffix = " - New Best!" if is_best else ""
+                    print(
+                        f"Step {step}: "
+                        f"train_loss={avg_train_loss:.5f}, "
+                        f"val_loss={val_loss:.5f}, "
+                        f"val_cratio={val_cratio:.5f}, "
+                        f"val_score={val_score:.5f}"
+                        f"{suffix}"
+                    )
+
+                    train_losses = []
+                    self.model.train()
+
+    def train_step(self, x, y, fg_mask):
+        # Forward
+        _, loss = self.forward_pass(x, y, fg_mask, use_amp=self.use_amp)
+
+        # Backward
+        self.optimizer.zero_grad()
+        self.scaler.scale(loss).backward()
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        return loss.item()
+
+    def validate(self, val_dataloader, step):
         """
-        Performs a single training epoch over the provided DataLoader.
-
-        Parameters
-        ----------
-        train_dataloader : torch.utils.data.DataLoader
-            DataLoader for the training dataset.
-        epoch : int
-            Current training epoch.
-
-        Returns
-        -------
-        loss : float
-            Average loss over the training epoch.
-        """
-        losses = list()
-        self.model.train()
-        for x, y, fg_mask in train_dataloader:
-            # Forward pass
-            hat_y, loss = self.forward_pass(
-                x, y, fg_mask, use_amp=self.use_amp
-            )
-
-            # Backward pass (loss-scaled for AMP stability)
-            self.optimizer.zero_grad()
-            self.scaler.scale(loss).backward()
-            self.scaler.step(self.optimizer)
-            self.scaler.update()
-
-            # Store loss for tensorboard
-            losses.append(float(loss.detach().cpu()))
-        self.writer.add_scalar("train_loss", np.mean(losses), epoch)
-        return np.mean(losses)
-
-    def validate_step(self, val_dataloader, epoch):
-        """
-        Validates the model over the provided DataLoader.
+        Validates the model over the validation dataset.
 
         Parameters
         ----------
         val_dataloader : torch.utils.data.DataLoader
             DataLoader for the validation dataset.
         epoch : int
-            Current training epoch.
+            Current number of gradient steps.
 
         Returns
         -------
@@ -273,10 +274,11 @@ def validate_step(self, val_dataloader, epoch):
         if len(val_dataloader.dataset) == 0:
             return float("nan"), float("nan"), False
 
+        # Run model over validation dataset
         losses = list()
         cratios = list()
         metric_rows = list()
-        with torch.no_grad():
+        with torch.inference_mode():
             self.model.eval()
             for x, y, raw, fg_mask in val_dataloader:
                 # Run model
@@ -285,7 +287,7 @@ def validate_step(self, val_dataloader, epoch):
                 )
 
                 # Evaluate result
-                losses.append(loss.detach().cpu())
+                losses.append(loss.item())
                 cratios.extend(self.compute_cratios(hat_y))
                 metric_rows.extend(
                     self.compute_metrics(hat_y, y, raw, fg_mask)
@@ -301,23 +303,13 @@ def validate_step(self, val_dataloader, epoch):
         score = checkpoint_score(agg, cratio, self.checkpoint_weights)
 
         # Log results
-        self.writer.add_scalar("val_loss", loss, epoch)
-        self.writer.add_scalar("val_cratio", cratio, epoch)
-        self.writer.add_scalar("val_score", score, epoch)
+        self.writer.add_scalar("val_loss", loss, step)
+        self.writer.add_scalar("val_cratio", cratio, step)
+        self.writer.add_scalar("val_score", score, step)
         for name, value in agg.items():
-            self.writer.add_scalar(f"val_{name}", value, epoch)
-
-        # Save every validated checkpoint so the best can be chosen offline.
-        # Skip epoch 0: the untrained net emits a near-constant, trivially
-        # compressible volume whose cratio-weighted score would beat every
-        # trained checkpoint despite its high loss. is_best is tracked only for
-        # the "New Best!" log line.
-        is_best = epoch > 0 and score < self.best_score
-        if is_best:
-            self.best_score = score
-        if epoch > 0:
-            self.save_model(epoch, score)
-        return loss, cratio, is_best
+            self.writer.add_scalar(f"val_{name}", value, step)
+
+        return loss, cratio, score
 
     def forward_pass(self, x, y, fg_mask, use_amp=None):
         """

From a6a374aae6b060274cfa68bfaaba0f8e2e595d6d Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Wed, 22 Jul 2026 21:37:59 +0000
Subject: [PATCH 02/10] refactor: simplified train

---
 .../machine_learning/train.py                 | 84 ++++++++++---------
 1 file changed, 43 insertions(+), 41 deletions(-)

diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index 6d5adfd..5921028 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -31,33 +31,6 @@
 from aind_exaspim_image_compression.utils import img_util, util
 
 
-def save_mip_png(path, img, low_pct=1.0, high_pct=99.9):
-    """
-    Writes a 3D volume as a contrast-stretched 8-bit PNG for easy viewing.
-
-    The volume is reduced to 2D with a maximum-intensity projection along the
-    z-axis (axis 0), then percentile-normalized to uint8 so that dim neurites
-    are visible in a standard image viewer.
-
-    Parameters
-    ----------
-    path : str
-        Output path for the PNG file.
-    img : numpy.ndarray
-        3D image volume with shape (D, H, W) in raw counts.
-    low_pct : float, optional
-        Lower percentile mapped to black. The default is 1.0.
-    high_pct : float, optional
-        Upper percentile mapped to white. The default is 99.9.
-    """
-    mip = img.max(axis=0).astype(np.float32)
-    lo, hi = np.percentile(mip, (low_pct, high_pct))
-    if hi <= lo:
-        hi = lo + 1.0
-    mip = np.clip((mip - lo) / (hi - lo), 0.0, 1.0)
-    io.imsave(path, np.rint(mip * 255).astype(np.uint8), check_contrast=False)
-
-
 class Trainer:
 
     def __init__(
@@ -146,7 +119,7 @@ def __init__(
         self.scaler = torch.amp.GradScaler("cuda", enabled=self.use_amp)
 
     # --- Core Routines ---
-    def run(self, train_dataset, val_dataset):
+    def __call__(self, train_dataset, val_dataset):
         """
         Runs the full training and validation loop.
 
@@ -193,18 +166,18 @@ def run(self, train_dataset, val_dataset):
         # Main
         step = 0
         best_score = np.inf
+        running_loss = 0
+        running_steps = 0
         self.model.train()
         for epoch in range(self.max_epochs):
-            # Epoch initializations
-            train_losses = []
             train_dataloader.set_epoch(epoch)
-
-            # Run train/val epoch
             for x, y, fg_mask in train_dataloader:
                 # Train
                 train_loss = self.train_step(x, y, fg_mask)
-                train_losses.append(train_loss)
                 self.scheduler.step()
+
+                running_loss += train_loss
+                running_steps += 1
                 step += 1
 
                 # Validate (if applicable)
@@ -213,15 +186,15 @@ def run(self, train_dataset, val_dataset):
                         val_dataloader,
                         step,
                     )
-                    is_best = epoch > 0 and val_score < best_score
-                    if is_best:
+                    is_best = val_score < best_score
+                    if epoch > 0 and is_best:
                         best_score = val_score
                     if epoch > 0:
                         self.save_model(step, val_score)
 
                 # Report results (if applicable)
                 if step % self.val_every == 0:
-                    avg_train_loss = np.mean(train_losses)
+                    avg_train_loss = running_loss / running_steps
                     self.writer.add_scalar(
                         "train_loss", avg_train_loss, step
                     )
@@ -236,7 +209,8 @@ def run(self, train_dataset, val_dataset):
                         f"{suffix}"
                     )
 
-                    train_losses = []
+                    running_loss = 0
+                    running_steps = 0
                     self.model.train()
 
     def train_step(self, x, y, fg_mask):
@@ -353,7 +327,7 @@ def compute_cratios(self, imgs):
         cratios = list()
         imgs = np.array(imgs.detach().cpu())
         for i in range(imgs.shape[0]):
-            img = self.transform.inverse(imgs[i, 0, ...])
+            img = self.transform.inverse(imgs[i, 0])
             cratios.append(img_util.compute_cratio(img, self.codec))
             if i < 10:
                 save_mip_png(f"{i}.png", img)
@@ -385,11 +359,11 @@ def compute_metrics(self, hat_y, y, raw, fg_mask):
         raws = np.array(raw.detach().cpu())
         masks = np.array(fg_mask.detach().cpu())
         for i in range(preds.shape[0]):
-            pred = self.transform.inverse(preds[i, 0, ...])
-            target = self.transform.inverse(targets[i, 0, ...])
+            pred = self.transform.inverse(preds[i, 0])
+            target = self.transform.inverse(targets[i, 0])
             rows.append(
                 evaluate_example(
-                    pred, raws[i, 0, ...], target, masks[i, 0, ...] > 0.5
+                    pred, raws[i, 0], target, masks[i, 0] > 0.5
                 )
             )
         return rows
@@ -478,3 +452,31 @@ def save_model(self, epoch, score):
             },
             path,
         )
+
+
+# --- Helpers ---
+def save_mip_png(path, img, low_pct=1.0, high_pct=99.9):
+    """
+    Writes a 3D volume as a contrast-stretched 8-bit PNG for easy viewing.
+
+    The volume is reduced to 2D with a maximum-intensity projection along the
+    z-axis (axis 0), then percentile-normalized to uint8 so that dim neurites
+    are visible in a standard image viewer.
+
+    Parameters
+    ----------
+    path : str
+        Output path for the PNG file.
+    img : numpy.ndarray
+        3D image volume with shape (D, H, W) in raw counts.
+    low_pct : float, optional
+        Lower percentile mapped to black. The default is 1.0.
+    high_pct : float, optional
+        Upper percentile mapped to white. The default is 99.9.
+    """
+    mip = img.max(axis=0).astype(np.float32)
+    lo, hi = np.percentile(mip, (low_pct, high_pct))
+    if hi <= lo:
+        hi = lo + 1.0
+    mip = np.clip((mip - lo) / (hi - lo), 0.0, 1.0)
+    io.imsave(path, np.rint(mip * 255).astype(np.uint8), check_contrast=False)

From 8abe6069131567cd5691a24d8c9506967aa6012a Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Wed, 22 Jul 2026 21:59:45 +0000
Subject: [PATCH 03/10] refactor: moved scheduler to main

---
 .../machine_learning/train.py                 | 40 +++++++++----------
 1 file changed, 18 insertions(+), 22 deletions(-)

diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index 5921028..94447e5 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -107,10 +107,6 @@ def __init__(
         self.model = model.to(device) if model else UNet().to(device)
         self._resume_transform_cfg = None
         self.optimizer = optim.AdamW(self.model.parameters(), lr=lr)
-        # T_max spans the whole run so the cosine anneals once. With a small
-        # T_max the LR returns to its peak every 2*T_max epochs, and each
-        # return destabilized training (growing periodic loss spikes).
-        self.scheduler = CosineAnnealingLR(self.optimizer, T_max=self.val_every)
         self.writer = SummaryWriter(log_dir=log_dir)
 
         # Scale the loss before backward so small float16 gradients do not
@@ -130,21 +126,10 @@ def __call__(self, train_dataset, val_dataset):
         val_dataset : ValidateDataset
             Dataset used for validation.
         """
-        # Initializations
-        print("Experiment:", os.path.basename(os.path.normpath(self.log_dir)))
+        # Create dataloaders
         if self._resume_transform_cfg is not None:
-            train_cfg = getattr(train_dataset.transform, "cfg", None)
-            val_cfg = getattr(val_dataset.transform, "cfg", None)
-            if train_cfg != self._resume_transform_cfg:
-                raise ValueError(
-                    "resume checkpoint transform does not match the training "
-                    "dataset transform"
-                )
-            if val_cfg != self._resume_transform_cfg:
-                raise ValueError(
-                    "resume checkpoint transform does not match the validation "
-                    "dataset transform"
-                )
+            self.check_transform_cfg(train_dataset, "train")
+            self.check_transform_cfg(val_dataset, "val")
 
         self.transform = train_dataset.transform
         train_dataloader = DataLoader(
@@ -163,7 +148,11 @@ def __call__(self, train_dataset, val_dataset):
             shuffle=False,
         )
 
-        # Main
+        # Create learning rate scheduler
+        total_steps = self.max_epochs * len(train_dataloader)
+        scheduler = CosineAnnealingLR(self.optimizer, T_max=total_steps)
+
+        # Training loop
         step = 0
         best_score = np.inf
         running_loss = 0
@@ -174,7 +163,7 @@ def __call__(self, train_dataset, val_dataset):
             for x, y, fg_mask in train_dataloader:
                 # Train
                 train_loss = self.train_step(x, y, fg_mask)
-                self.scheduler.step()
+                scheduler.step()
 
                 running_loss += train_loss
                 running_steps += 1
@@ -183,8 +172,7 @@ def __call__(self, train_dataset, val_dataset):
                 # Validate (if applicable)
                 if step % self.val_every == 0:
                     val_loss, val_cratio, val_score = self.validate(
-                        val_dataloader,
-                        step,
+                        val_dataloader, step,
                     )
                     is_best = val_score < best_score
                     if epoch > 0 and is_best:
@@ -323,6 +311,14 @@ def forward_pass(self, x, y, fg_mask, use_amp=None):
             return hat_y, loss
 
     # --- Helpers ---
+    def check_transform_cfg(self, dataset, dataset_name):
+        cfg = getattr(dataset.transform, "cfg", None)
+        if cfg != self._resume_transform_cfg:
+            raise ValueError(
+                f"Resume checkpoint transform does not match the "
+                f"{dataset_name} dataset transform."
+            )
+
     def compute_cratios(self, imgs):
         cratios = list()
         imgs = np.array(imgs.detach().cpu())

From 90a1e49bdf76c0e9135aba0529dedd0f7228b7e5 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Wed, 22 Jul 2026 22:03:55 +0000
Subject: [PATCH 04/10] refactor: switch order of scheduler and optimizer

---
 src/aind_exaspim_image_compression/machine_learning/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index 94447e5..e53ba4e 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -162,8 +162,8 @@ def __call__(self, train_dataset, val_dataset):
             train_dataloader.set_epoch(epoch)
             for x, y, fg_mask in train_dataloader:
                 # Train
-                train_loss = self.train_step(x, y, fg_mask)
                 scheduler.step()
+                train_loss = self.train_step(x, y, fg_mask)
 
                 running_loss += train_loss
                 running_steps += 1

From fd2e81628b0dd0f5d5483ced28953a2c479579b3 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Wed, 22 Jul 2026 22:08:09 +0000
Subject: [PATCH 05/10] refactor: report total steps

---
 src/aind_exaspim_image_compression/machine_learning/train.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index e53ba4e..49a98df 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -151,6 +151,7 @@ def __call__(self, train_dataset, val_dataset):
         # Create learning rate scheduler
         total_steps = self.max_epochs * len(train_dataloader)
         scheduler = CosineAnnealingLR(self.optimizer, T_max=total_steps)
+        print("Total Optimizer Steps:", total_steps)
 
         # Training loop
         step = 0
@@ -187,7 +188,7 @@ def __call__(self, train_dataset, val_dataset):
                         "train_loss", avg_train_loss, step
                     )
 
-                    suffix = " - New Best!" if is_best else ""
+                    suffix = " - New Best!" if epoch > 0 and is_best else ""
                     print(
                         f"Step {step}: "
                         f"train_loss={avg_train_loss:.5f}, "

From 7bc3fb1c0aa055cb4eca10dbb6fcc31653466bce Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Wed, 22 Jul 2026 22:19:23 +0000
Subject: [PATCH 06/10] bug: dataloader length

---
 src/aind_exaspim_image_compression/machine_learning/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index 49a98df..2d25d5c 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -149,7 +149,7 @@ def __call__(self, train_dataset, val_dataset):
         )
 
         # Create learning rate scheduler
-        total_steps = self.max_epochs * len(train_dataloader)
+        total_steps = self.max_epochs * len(train_dataset) // self.batch_size
         scheduler = CosineAnnealingLR(self.optimizer, T_max=total_steps)
         print("Total Optimizer Steps:", total_steps)
 

From f5b35d4dcf68cec06c650ba1887c68b00b5977e0 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Wed, 22 Jul 2026 22:30:25 +0000
Subject: [PATCH 07/10] refactor: switcher scheduler and opt

---
 src/aind_exaspim_image_compression/machine_learning/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index 2d25d5c..f34e39e 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -163,8 +163,8 @@ def __call__(self, train_dataset, val_dataset):
             train_dataloader.set_epoch(epoch)
             for x, y, fg_mask in train_dataloader:
                 # Train
-                scheduler.step()
                 train_loss = self.train_step(x, y, fg_mask)
+                scheduler.step()
 
                 running_loss += train_loss
                 running_steps += 1

From 4751a764cc7c20409153f0a1c971ec825a0c592b Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Thu, 23 Jul 2026 17:47:40 +0000
Subject: [PATCH 08/10] refactor: suggested edits

---
 scripts/train_bm4dnet.py                      |  2 +-
 .../machine_learning/train.py                 | 45 +++++++++++++------
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/scripts/train_bm4dnet.py b/scripts/train_bm4dnet.py
index 15f8335..a537505 100644
--- a/scripts/train_bm4dnet.py
+++ b/scripts/train_bm4dnet.py
@@ -141,7 +141,7 @@ def train(train_cache_dir, val_cache_dir):
     use_amp = False
     use_amp_validation = False
     # Validate (and consider a checkpoint) after every full-cache epoch.
-    val_every = 1
+    val_every = 1000
     seed = 42
 
     # Signal-preserving loss + target/sampling (Parts E/F). preserve_foreground
diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index f34e39e..f332af5 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -13,6 +13,7 @@
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torch.utils.tensorboard import SummaryWriter
 
+import math
 import numpy as np
 import os
 import torch
@@ -115,7 +116,7 @@ def __init__(
         self.scaler = torch.amp.GradScaler("cuda", enabled=self.use_amp)
 
     # --- Core Routines ---
-    def __call__(self, train_dataset, val_dataset):
+    def run(self, train_dataset, val_dataset):
         """
         Runs the full training and validation loop.
 
@@ -149,9 +150,12 @@ def __call__(self, train_dataset, val_dataset):
         )
 
         # Create learning rate scheduler
-        total_steps = self.max_epochs * len(train_dataset) // self.batch_size
+        total_steps = self.max_epochs * math.ceil(len(train_dataset) / self.batch_size)
+        num_validations = math.ceil(total_steps / self.val_every)
         scheduler = CosineAnnealingLR(self.optimizer, T_max=total_steps)
-        print("Total Optimizer Steps:", total_steps)
+        print(f"Total optimizer steps: {total_steps}")
+        print(f"Validation interval: every {self.val_every} steps")
+        print(f"Expected validations: {num_validations}")
 
         # Training loop
         step = 0
@@ -160,6 +164,7 @@ def __call__(self, train_dataset, val_dataset):
         running_steps = 0
         self.model.train()
         for epoch in range(self.max_epochs):
+            print(f"Starting epoch {epoch} / {self.max_epochs}...")
             train_dataloader.set_epoch(epoch)
             for x, y, fg_mask in train_dataloader:
                 # Train
@@ -175,11 +180,13 @@ def __call__(self, train_dataset, val_dataset):
                     val_loss, val_cratio, val_score = self.validate(
                         val_dataloader, step,
                     )
-                    is_best = val_score < best_score
-                    if epoch > 0 and is_best:
-                        best_score = val_score
-                    if epoch > 0:
-                        self.save_model(step, val_score)
+                    self.save_model(step, val_score)
+
+                    if val_score is not None:
+                        is_best = val_score < best_score
+                        best_score = min(best_score, val_score)
+                    else:
+                        is_best = False
 
                 # Report results (if applicable)
                 if step % self.val_every == 0:
@@ -188,7 +195,7 @@ def __call__(self, train_dataset, val_dataset):
                         "train_loss", avg_train_loss, step
                     )
 
-                    suffix = " - New Best!" if epoch > 0 and is_best else ""
+                    suffix = " - New Best!" if is_best else ""
                     print(
                         f"Step {step}: "
                         f"train_loss={avg_train_loss:.5f}, "
@@ -202,6 +209,18 @@ def __call__(self, train_dataset, val_dataset):
                     running_steps = 0
                     self.model.train()
 
+            # Final model validation
+            val_loss, val_cratio, val_score = self.validate(
+                val_dataloader, step,
+            )
+            self.save_model(step, val_score)
+            print(
+                "Final model validation:"
+                f"val_loss={val_loss:.5f}, "
+                f"val_cratio={val_cratio:.5f}, "
+                f"val_score={val_score:.5f}"
+            )
+
     def train_step(self, x, y, fg_mask):
         # Forward
         _, loss = self.forward_pass(x, y, fg_mask, use_amp=self.use_amp)
@@ -221,7 +240,7 @@ def validate(self, val_dataloader, step):
         ----------
         val_dataloader : torch.utils.data.DataLoader
             DataLoader for the validation dataset.
-        epoch : int
+        step : int
             Current number of gradient steps.
 
         Returns
@@ -230,12 +249,12 @@ def validate(self, val_dataloader, step):
             Average loss over the validation dataset.
         cratio : float
             Average compression ratio over the validation dataset.
-        is_best : bool
-            Indication of whether the model is the best so far.
+        score : float or None
+            Model score on validation dataset.
         """
         # Skip if there are no validation examples
         if len(val_dataloader.dataset) == 0:
-            return float("nan"), float("nan"), False
+            return float("nan"), float("nan"), None
 
         # Run model over validation dataset
         losses = list()

From 085637b740fcf47340fb2da272e379db5a5e3bc5 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Thu, 23 Jul 2026 20:36:53 +0000
Subject: [PATCH 09/10] refactor: review fixes

---
 .../machine_learning/data_handling.py         |   4 +
 .../machine_learning/train.py                 | 115 ++++++++----------
 2 files changed, 56 insertions(+), 63 deletions(-)

diff --git a/src/aind_exaspim_image_compression/machine_learning/data_handling.py b/src/aind_exaspim_image_compression/machine_learning/data_handling.py
index 242131b..eb7691f 100644
--- a/src/aind_exaspim_image_compression/machine_learning/data_handling.py
+++ b/src/aind_exaspim_image_compression/machine_learning/data_handling.py
@@ -19,6 +19,7 @@
 
 import logging
 import fastremap
+import math
 import numpy as np
 import os
 import queue
@@ -1314,6 +1315,9 @@ def _load_batch(self, executor, indices):
                 batched[j][i, 0, ...] = field
         return tuple(to_tensor(arr) for arr in batched)
 
+    def __len__(self):
+        return math.ceil(len(self.dataset) / self.batch_size)
+
 
 # --- Helpers ---
 def init_datasets(
diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index f332af5..3ed0fd7 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -13,7 +13,6 @@
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torch.utils.tensorboard import SummaryWriter
 
-import math
 import numpy as np
 import os
 import torch
@@ -43,7 +42,7 @@ def __init__(
         max_epochs=400,
         model=None,
         use_amp=True,
-        use_amp_validation=False,
+        use_amp_val=False,
         checkpoint_weights=None,
         fg_weight=20.0,
         num_workers=None,
@@ -71,7 +70,7 @@ def __init__(
         use_amp : bool, optional
             Whether to use CUDA float16 autocast and gradient scaling during
             training. Default is True.
-        use_amp_validation : bool, optional
+        use_amp_val : bool, optional
             Whether to use CUDA float16 autocast during validation and
             checkpoint selection. Default is False so validation matches FP32
             production inference.
@@ -99,8 +98,8 @@ def __init__(
         self.prefetch = prefetch
         self.val_every = max(1, int(val_every))
         self.seed = seed
-        self.use_amp = bool(use_amp)
-        self.use_amp_validation = bool(use_amp_validation)
+        self.use_amp = bool(use_amp and device.startswith("cuda"))
+        self.use_amp_val = bool(use_amp_val and device.startswith("cuda"))
 
         self.codec = blosc.Blosc(cname="zstd", clevel=6, shuffle=blosc.SHUFFLE)
         self.criterion = SignalPreservingLoss(fg_weight=fg_weight)
@@ -150,17 +149,16 @@ def run(self, train_dataset, val_dataset):
         )
 
         # Create learning rate scheduler
-        total_steps = self.max_epochs * math.ceil(len(train_dataset) / self.batch_size)
-        num_validations = math.ceil(total_steps / self.val_every)
+        total_steps = self.max_epochs * len(train_dataloader)
+        expected_validations = total_steps // self.val_every
         scheduler = CosineAnnealingLR(self.optimizer, T_max=total_steps)
         print(f"Total optimizer steps: {total_steps}")
         print(f"Validation interval: every {self.val_every} steps")
-        print(f"Expected validations: {num_validations}")
+        print(f"Expected validations: {expected_validations}")
 
         # Training loop
         step = 0
-        best_score = np.inf
-        running_loss = 0
+        running_loss = 0.0
         running_steps = 0
         self.model.train()
         for epoch in range(self.max_epochs):
@@ -168,58 +166,31 @@ def run(self, train_dataset, val_dataset):
             train_dataloader.set_epoch(epoch)
             for x, y, fg_mask in train_dataloader:
                 # Train
-                train_loss = self.train_step(x, y, fg_mask)
+                loss = self.train_step(x, y, fg_mask)
                 scheduler.step()
 
-                running_loss += train_loss
+                running_loss += loss
                 running_steps += 1
                 step += 1
 
                 # Validate (if applicable)
                 if step % self.val_every == 0:
-                    val_loss, val_cratio, val_score = self.validate(
-                        val_dataloader, step,
-                    )
-                    self.save_model(step, val_score)
-
-                    if val_score is not None:
-                        is_best = val_score < best_score
-                        best_score = min(best_score, val_score)
-                    else:
-                        is_best = False
-
-                # Report results (if applicable)
-                if step % self.val_every == 0:
-                    avg_train_loss = running_loss / running_steps
-                    self.writer.add_scalar(
-                        "train_loss", avg_train_loss, step
-                    )
-
-                    suffix = " - New Best!" if is_best else ""
-                    print(
-                        f"Step {step}: "
-                        f"train_loss={avg_train_loss:.5f}, "
-                        f"val_loss={val_loss:.5f}, "
-                        f"val_cratio={val_cratio:.5f}, "
-                        f"val_score={val_score:.5f}"
-                        f"{suffix}"
-                    )
+                    # Summarize train progress
+                    avg_loss = running_loss / running_steps
+                    label = f"Step {step}:  train_loss={avg_loss:.5f}, "
+                    self.writer.add_scalar("train_loss", avg_loss, step)
+
+                    # Call validation
+                    self.validate_and_checkpoint(val_dataloader, step, label)
 
+                    # Reset counters for training session
                     running_loss = 0
                     running_steps = 0
                     self.model.train()
 
-            # Final model validation
-            val_loss, val_cratio, val_score = self.validate(
-                val_dataloader, step,
-            )
-            self.save_model(step, val_score)
-            print(
-                "Final model validation:"
-                f"val_loss={val_loss:.5f}, "
-                f"val_cratio={val_cratio:.5f}, "
-                f"val_score={val_score:.5f}"
-            )
+        # Final model validation
+        if step % self.val_every != 0:
+            self.validate_and_checkpoint(val_dataloader, step, "Final: ")
 
     def train_step(self, x, y, fg_mask):
         # Forward
@@ -232,6 +203,23 @@ def train_step(self, x, y, fg_mask):
         self.scaler.update()
         return loss.item()
 
+    def validate_and_checkpoint(self, val_dataloader, step, label):
+        # Validate
+        loss, cratio, score = self.validate(val_dataloader, step)
+        if score is None:
+            print(f"{label}: validation skipped (empty validation dataset)")
+            return False
+
+        # Save checkpoint and report reuslts
+        self.save_model(step, score)
+        print(
+            f"{label}"
+            f"val_loss={loss:.5f}, "
+            f"val_cratio={cratio:.5f}, "
+            f"val_score={score:.5f}"
+        )
+        return True
+
     def validate(self, val_dataloader, step):
         """
         Validates the model over the validation dataset.
@@ -260,19 +248,19 @@ def validate(self, val_dataloader, step):
         losses = list()
         cratios = list()
         metric_rows = list()
+        self.model.eval()
         with torch.inference_mode():
-            self.model.eval()
             for x, y, raw, fg_mask in val_dataloader:
                 # Run model
-                hat_y, loss = self.forward_pass(
-                    x, y, fg_mask, use_amp=self.use_amp_validation
+                y_pred, loss = self.forward_pass(
+                    x, y, fg_mask, use_amp=self.use_amp_val
                 )
 
                 # Evaluate result
                 losses.append(loss.item())
-                cratios.extend(self.compute_cratios(hat_y))
+                cratios.extend(self.compute_cratios(y_pred))
                 metric_rows.extend(
-                    self.compute_metrics(hat_y, y, raw, fg_mask)
+                    self.compute_metrics(y_pred, y, raw, fg_mask)
                 )
 
         # Aggregate results
@@ -311,7 +299,7 @@ def forward_pass(self, x, y, fg_mask, use_amp=None):
 
         Returns
         -------
-        hat_y : torch.Tensor
+        y_pred : torch.Tensor
             Model predictions.
         loss : torch.Tensor
             Computed loss value.
@@ -326,9 +314,9 @@ def forward_pass(self, x, y, fg_mask, use_amp=None):
             x = x.to(self.device)
             y = y.to(self.device)
             fg_mask = fg_mask.to(self.device)
-            hat_y = self.model(x)
-            loss = self.criterion(hat_y, y, fg_mask)
-            return hat_y, loss
+            y_pred = self.model(x)
+            loss = self.criterion(y_pred, y, fg_mask)
+            return y_pred, loss
 
     # --- Helpers ---
     def check_transform_cfg(self, dataset, dataset_name):
@@ -349,13 +337,13 @@ def compute_cratios(self, imgs):
                 save_mip_png(f"{i}.png", img)
         return cratios
 
-    def compute_metrics(self, hat_y, y, raw, fg_mask):
+    def compute_metrics(self, y_pred, y, raw, fg_mask):
         """
         Computes per-example neurite-preservation metrics in count space.
 
         Parameters
         ----------
-        hat_y : torch.Tensor
+        y_pred : torch.Tensor
             Model predictions in the normalized transform domain.
         y : torch.Tensor
             BM4D targets in the normalized transform domain.
@@ -370,7 +358,7 @@ def compute_metrics(self, hat_y, y, raw, fg_mask):
             One metric dictionary per example in the batch.
         """
         rows = list()
-        preds = np.array(hat_y.detach().cpu())
+        preds = np.array(y_pred.detach().cpu())
         targets = np.array(y.detach().cpu())
         raws = np.array(raw.detach().cpu())
         masks = np.array(fg_mask.detach().cpu())
@@ -435,7 +423,7 @@ def save_config(self, config):
             "val_every": self.val_every,
             "seed": self.seed,
             "use_amp": self.use_amp,
-            "use_amp_validation": self.use_amp_validation,
+            "use_amp_val": self.use_amp_val,
             "fg_weight": getattr(self.criterion, "fg_weight", None),
             "checkpoint_weights": self.checkpoint_weights,
             "lr": self.optimizer.param_groups[0]["lr"],
@@ -458,6 +446,7 @@ def save_model(self, epoch, score):
             embedded in the filename so checkpoints can be ranked offline.
         """
         date = datetime.today().strftime("%Y%m%d")
+        score = score or np.inf
         filename = f"BM4DNet-{date}-{epoch}-{score:.6f}.pth"
         path = os.path.join(self.log_dir, filename)
         torch.save(

From 1bd413b99e28d1f6a4cbe3f83fa2cd3611770926 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Thu, 23 Jul 2026 21:03:26 +0000
Subject: [PATCH 10/10] refactor: save model bug, expected vals bug

---
 .../machine_learning/train.py                     | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py
index 3ed0fd7..a473e74 100644
--- a/src/aind_exaspim_image_compression/machine_learning/train.py
+++ b/src/aind_exaspim_image_compression/machine_learning/train.py
@@ -13,6 +13,7 @@
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torch.utils.tensorboard import SummaryWriter
 
+import math
 import numpy as np
 import os
 import torch
@@ -42,7 +43,7 @@ def __init__(
         max_epochs=400,
         model=None,
         use_amp=True,
-        use_amp_val=False,
+        use_amp_validation=False,
         checkpoint_weights=None,
         fg_weight=20.0,
         num_workers=None,
@@ -70,7 +71,7 @@ def __init__(
         use_amp : bool, optional
             Whether to use CUDA float16 autocast and gradient scaling during
             training. Default is True.
-        use_amp_val : bool, optional
+        use_amp_validation : bool, optional
             Whether to use CUDA float16 autocast during validation and
             checkpoint selection. Default is False so validation matches FP32
             production inference.
@@ -99,7 +100,7 @@ def __init__(
         self.val_every = max(1, int(val_every))
         self.seed = seed
         self.use_amp = bool(use_amp and device.startswith("cuda"))
-        self.use_amp_val = bool(use_amp_val and device.startswith("cuda"))
+        self.use_amp_validation = bool(use_amp_validation and device.startswith("cuda"))
 
         self.codec = blosc.Blosc(cname="zstd", clevel=6, shuffle=blosc.SHUFFLE)
         self.criterion = SignalPreservingLoss(fg_weight=fg_weight)
@@ -150,7 +151,7 @@ def run(self, train_dataset, val_dataset):
 
         # Create learning rate scheduler
         total_steps = self.max_epochs * len(train_dataloader)
-        expected_validations = total_steps // self.val_every
+        expected_validations = math.ceil(total_steps / self.val_every)
         scheduler = CosineAnnealingLR(self.optimizer, T_max=total_steps)
         print(f"Total optimizer steps: {total_steps}")
         print(f"Validation interval: every {self.val_every} steps")
@@ -253,7 +254,7 @@ def validate(self, val_dataloader, step):
             for x, y, raw, fg_mask in val_dataloader:
                 # Run model
                 y_pred, loss = self.forward_pass(
-                    x, y, fg_mask, use_amp=self.use_amp_val
+                    x, y, fg_mask, use_amp=self.use_amp_validation
                 )
 
                 # Evaluate result
@@ -423,7 +424,7 @@ def save_config(self, config):
             "val_every": self.val_every,
             "seed": self.seed,
             "use_amp": self.use_amp,
-            "use_amp_val": self.use_amp_val,
+            "use_amp_validation": self.use_amp_validation,
             "fg_weight": getattr(self.criterion, "fg_weight", None),
             "checkpoint_weights": self.checkpoint_weights,
             "lr": self.optimizer.param_groups[0]["lr"],
@@ -446,7 +447,7 @@ def save_model(self, epoch, score):
             embedded in the filename so checkpoints can be ranked offline.
         """
         date = datetime.today().strftime("%Y%m%d")
-        score = score or np.inf
+        score = np.inf if score is None else score
         filename = f"BM4DNet-{date}-{epoch}-{score:.6f}.pth"
         path = os.path.join(self.log_dir, filename)
         torch.save(