Add option to disable AMP during training/validation

- enabling AMP can decrease compression ratios during inference
- all current models were trained using AMP. We need to be able to test training without.

Author: carshadi

scripts/train_bm4dnet.py

@@ -81,6 +81,8 @@ def train(train_cache_dir, val_cache_dir):
         lr=lr,
         max_epochs=max_epochs,
         model=model,
+        use_amp=use_amp,
+        use_amp_validation=use_amp_validation,
         fg_weight=fg_weight,
         checkpoint_weights=checkpoint_weights,
         num_workers=0,
@@ -97,6 +99,8 @@ def train(train_cache_dir, val_cache_dir):
             "resume_path": resume_path,
             "transform_cfg": transform.cfg,
             "preserve_foreground": preserve_foreground,
+            "use_amp": use_amp,
+            "use_amp_validation": use_amp_validation,
         }
     )
 
@@ -131,7 +135,11 @@ def train(train_cache_dir, val_cache_dir):
     # Training parameters
     batch_size = 32
     lr = 1e-3
-    max_epochs = 30
+    max_epochs = 50
+    # Use AMP for optimization, but keep validation/checkpoint selection in
+    # FP32 so its cratio matches production FP32 inference.
+    use_amp = False
+    use_amp_validation = False
     # Validate (and consider a checkpoint) after every full-cache epoch.
     val_every = 1
     seed = 42

src/aind_exaspim_image_compression/machine_learning/train.py

@@ -8,7 +8,6 @@
 
 """
 
-from contextlib import nullcontext
 from datetime import datetime
 from numcodecs import blosc
 from torch.optim.lr_scheduler import CosineAnnealingLR
@@ -71,6 +70,7 @@ def __init__(
         max_epochs=400,
         model=None,
         use_amp=True,
+        use_amp_validation=False,
         checkpoint_weights=None,
         fg_weight=20.0,
         num_workers=None,
@@ -96,7 +96,12 @@ def __init__(
         model : None or nn.Module, optional
             Model to be trained on the given datasets. Default is None.
         use_amp : bool, optional
-            Indication of whether to use mixed precision. Default is True.
+            Whether to use CUDA float16 autocast and gradient scaling during
+            training. Default is True.
+        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.
         val_every : int, optional
             Run validation (and checkpoint selection) every this many epochs;
             the final epoch is always validated. The count-space metrics are
@@ -120,6 +125,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.codec = blosc.Blosc(cname="zstd", clevel=5, shuffle=blosc.SHUFFLE)
         self.criterion = SignalPreservingLoss(fg_weight=fg_weight)
@@ -134,15 +141,10 @@ def __init__(
         self.scheduler = CosineAnnealingLR(self.optimizer, T_max=max_epochs)
         self.writer = SummaryWriter(log_dir=log_dir)
 
-        if use_amp:
-            self.autocast = torch.autocast(device_type="cuda", dtype=torch.float16)
-        else:
-            self.autocast = nullcontext()
-
         # Scale the loss before backward so small float16 gradients do not
         # underflow (and are unscaled before the step). Disabled => no-op, so
         # the same code path is correct with and without AMP.
-        self.scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
+        self.scaler = torch.amp.GradScaler("cuda", enabled=self.use_amp)
 
     # --- Core Routines ---
     def run(self, train_dataset, val_dataset):
@@ -232,7 +234,9 @@ def train_step(self, train_dataloader, epoch):
         self.model.train()
         for x, y, fg_mask in train_dataloader:
             # Forward pass
-            hat_y, loss = self.forward_pass(x, y, fg_mask)
+            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()
@@ -276,7 +280,9 @@ def validate_step(self, val_dataloader, epoch):
             self.model.eval()
             for x, y, raw, fg_mask in val_dataloader:
                 # Run model
-                hat_y, loss = self.forward_pass(x, y, fg_mask)
+                hat_y, loss = self.forward_pass(
+                    x, y, fg_mask, use_amp=self.use_amp_validation
+                )
 
                 # Evaluate result
                 losses.append(loss.detach().cpu())
@@ -313,7 +319,7 @@ def validate_step(self, val_dataloader, epoch):
             self.save_model(epoch, score)
         return loss, cratio, is_best
 
-    def forward_pass(self, x, y, fg_mask):
+    def forward_pass(self, x, y, fg_mask, use_amp=None):
         """
         Performs a forward pass through the model and computes loss.
 
@@ -325,6 +331,9 @@ def forward_pass(self, x, y, fg_mask):
             Target tensor with shape (B, C, D, H, W).
         fg_mask : torch.Tensor
             Foreground mask (0/1) with shape (B, C, D, H, W).
+        use_amp : bool, optional
+            Whether to autocast this forward pass to float16. Defaults to the
+            training AMP setting.
 
         Returns
         -------
@@ -333,7 +342,13 @@ def forward_pass(self, x, y, fg_mask):
         loss : torch.Tensor
             Computed loss value.
         """
-        with self.autocast:
+        if use_amp is None:
+            use_amp = self.use_amp
+        with torch.autocast(
+            device_type=torch.device(self.device).type,
+            dtype=torch.float16,
+            enabled=bool(use_amp),
+        ):
             x = x.to(self.device)
             y = y.to(self.device)
             fg_mask = fg_mask.to(self.device)
@@ -437,6 +452,8 @@ def save_config(self, config):
             "prefetch": self.prefetch,
             "val_every": self.val_every,
             "seed": self.seed,
+            "use_amp": self.use_amp,
+            "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"],

tests/test_full_cache_training.py

@@ -4,6 +4,7 @@
 import os
 import tempfile
 import unittest
+from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import MagicMock, call, patch
 
@@ -118,6 +119,42 @@ def test_validation_is_ordered_and_final_partial_batch_is_retained(self):
 class TrainerShuffleTest(unittest.TestCase):
     """Tests Trainer wiring for shuffled training and ordered validation."""
 
+    def test_training_and_validation_amp_are_configured_separately(self):
+        """Training can use AMP while validation remains full precision."""
+        model = torch.nn.Linear(1, 1)
+        values = torch.zeros((1, 1, 1, 1, 1))
+        with tempfile.TemporaryDirectory() as directory:
+            trainer = Trainer(
+                directory,
+                device="cpu",
+                model=model,
+                max_epochs=1,
+                use_amp=True,
+                use_amp_validation=False,
+            )
+            try:
+                with patch(
+                    "aind_exaspim_image_compression.machine_learning."
+                    "train.torch.autocast",
+                    side_effect=[nullcontext(), nullcontext()],
+                ) as autocast:
+                    trainer.forward_pass(
+                        values, values, values, use_amp=trainer.use_amp
+                    )
+                    trainer.forward_pass(
+                        values,
+                        values,
+                        values,
+                        use_amp=trainer.use_amp_validation,
+                    )
+
+                self.assertTrue(autocast.call_args_list[0].kwargs["enabled"])
+                self.assertFalse(
+                    autocast.call_args_list[1].kwargs["enabled"]
+                )
+            finally:
+                trainer.writer.close()
+
     def test_trainer_sets_epoch_and_persists_seed(self):
         """Trainer passes its seed, sets every epoch, and saves the seed."""
         model = torch.nn.Linear(1, 1)
@@ -177,7 +214,10 @@ def test_trainer_sets_epoch_and_persists_seed(self):
                     os.path.join(trainer.log_dir, "config.json"),
                     encoding="utf-8",
                 ) as file:
-                    self.assertEqual(json.load(file)["seed"], 42)
+                    config = json.load(file)
+                    self.assertEqual(config["seed"], 42)
+                    self.assertFalse(config["use_amp"])
+                    self.assertFalse(config["use_amp_validation"])
             finally:
                 trainer.writer.close()
 

tests/test_train_bm4dnet.py

@@ -101,6 +101,8 @@ def test_training_uses_only_cached_datasets(self):
             "lr": 1e-3,
             "max_epochs": 3,
             "model": object(),
+            "use_amp": True,
+            "use_amp_validation": False,
             "fg_weight": 0,
             "checkpoint_weights": {"fg_mae": 1.0},
             "val_every": 2,
@@ -159,11 +161,17 @@ def test_training_uses_only_cached_datasets(self):
                 self.assertEqual(config["train_cache_dir"], str(train_cache))
                 self.assertEqual(config["val_cache_dir"], str(val_cache))
                 self.assertEqual(config["transform_cfg"], transform.cfg)
+                self.assertTrue(config["use_amp"])
+                self.assertFalse(config["use_amp_validation"])
                 self.assertNotIn("n_train_examples_per_epoch", config)
                 self.assertNotIn("brain_ids_path", config)
                 self.assertNotIn("sigma_bm4d", config)
                 trainer_call = trainer_factory.call_args
                 self.assertEqual(trainer_call.kwargs["seed"], 42)
+                self.assertTrue(trainer_call.kwargs["use_amp"])
+                self.assertFalse(
+                    trainer_call.kwargs["use_amp_validation"]
+                )
         finally:
             for key, value in previous.items():
                 if value is None: