Enable missing mixed precision plug for joint training setup (#1302)

* Enable missing mixed precision plug for joint training setup

* Pin per model joint training configuration

Author: anwai98

finetuning/v2/generalist/train_joint.py

@@ -4,13 +4,27 @@
 import torch
 
 
+CHOSEN_PARAMETERS = {
+    "hvit_t": (10, 10, 5),
+    "hvit_s": (10, 10, 5),
+    "hvit_b": (8, 10, 5),
+    "hvit_l": (8, 8, 4),
+}
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--n_epochs", type=int, default=100)
     parser.add_argument("--n_iterations", type=int, default=None)
+    parser.add_argument("--model_type", default="hvit_t", choices=["hvit_t", "hvit_s", "hvit_b", "hvit_l"])
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--dataset_choice", default="both", choices=["lm", "em", "both"])
     args = parser.parse_args()
 
-    model_type = "hvit_t"
+    model_type = args.model_type
+    # Pinned per-model config (batch_size_2d, z_slices, max_num_objects); not CLI-tunable.
+    batch_size_2d, z_slice, max_num_objects = CHOSEN_PARAMETERS[model_type]
+    z_slices = [z_slice]
     data_path = "/mnt/vast-nhr/projects/cidas/cca/data"
     save_root = os.environ.get("SAVE_ROOT", "/mnt/vast-nhr/projects/cidas/cca/models/micro_sam2/joint/v0")
 
@@ -23,17 +37,17 @@ def main():
         name=name,
         model_type=model_type,
         input_path=data_path,
-        batch_size=1,
-        batch_size_2d=8,
-        z_slices=[8],
-        dataset_choice="both",
+        batch_size=args.batch_size,
+        batch_size_2d=batch_size_2d,
+        z_slices=z_slices,
+        dataset_choice=args.dataset_choice,
         n_workers=8,
         n_epochs=args.n_epochs,
         n_iterations=args.n_iterations,
         lr=1e-5,  # single LR for all parameters
         save_root=save_root,
         checkpoint_path=None,  # downloads default SAM2 weights if None
-        max_num_objects=5,  # lower than interactive-only (8): joint also holds the UNETR decoder
+        max_num_objects=max_num_objects,  # lower than interactive-only (8): joint also holds the UNETR decoder
         prob_to_use_pt_input=1.0,  # always point/box prompts, never the GT mask
         prob_to_use_box_input=0.5,  # conditional prob of a box instead of a click
         num_frames_to_correct=2,  # max frames per volume receiving correction clicks
@@ -58,6 +72,17 @@ def main():
         from micro_sam.v2.training import train_joint_sam2
         train_joint_sam2(**common)
 
+    rank = int(os.environ.get("RANK", "0"))
+    if torch.cuda.is_available() and rank == 0:
+        peak_alloc = torch.cuda.max_memory_allocated() / 1024**3
+        peak_reserved = torch.cuda.max_memory_reserved() / 1024**3
+        print(
+            f"[peak-memory] model_type={model_type} batch_size={args.batch_size} "
+            f"batch_size_2d={batch_size_2d} z_slices={z_slices} "
+            f"max_num_objects={max_num_objects} "
+            f"allocated={peak_alloc:.2f}GiB reserved={peak_reserved:.2f}GiB", flush=True
+        )
+
 
 if __name__ == "__main__":
     main()

micro_sam/v2/training/sam2_trainer.py

@@ -128,6 +128,15 @@ def _amp_context(self):
             return torch.amp.autocast(device_type="cuda", dtype=self.amp_dtype)
         return contextlib.nullcontext()
 
+    def _train_epoch(self, progress):
+        # mixed_precision is disabled on the parent (see __init__), so DefaultTrainer would pass
+        # contextlib.nullcontext as forward_context. Inject our bf16 autocast instead so every
+        # _train_epoch_impl (this trainer and JointSam2Trainer) runs the forward pass under AMP.
+        return self._train_epoch_impl(progress, self._amp_context, self._sam2_backprop)
+
+    def _validate(self):
+        return self._validate_impl(self._amp_context)
+
     def _check_input_normalization(self, x, input_check_done):
         if not input_check_done:
             data_min, data_max = x.min(), x.max()
@@ -182,7 +191,7 @@ def _train_epoch_impl(self, progress, forward_context, backprop):
             input_check_done = self._check_input_normalization(x, input_check_done)
             self.optimizer.zero_grad()
 
-            with self._amp_context():
+            with forward_context():
                 loss, batch, outputs = self._interactive_step(x, y)
 
             grad_norm = self._sam2_backprop(loss)
@@ -235,7 +244,7 @@ def _validate_impl(self, forward_context):
             with torch.no_grad():
                 for i, (x, y) in enumerate(self.val_loader):
                     input_check_done = self._check_input_normalization(x, input_check_done)
-                    with self._amp_context():
+                    with forward_context():
                         loss, batch, outputs = self._interactive_step(x, y)
                         val_loss += loss.item()
                         val_dice_loss += self.metric(outputs, batch.masks).item()

micro_sam/v2/transforms/labels.py

@@ -132,7 +132,8 @@ def __call__(self, labels: np.ndarray) -> np.ndarray:
             labels = labels.byteswap().view(labels.dtype.newbyteorder())
 
         if self.apply_label:
-            labels = connected_components(labels).astype("uint32")
+            # Cast to uint32: connected_components rejects int16 and labels fit uint32.
+            labels = connected_components(labels.astype("uint32")).astype("uint32")
         else:  # Otherwise just relabel the segmentation.
             # Cast to uint32: relabel_sequential rejects uint8/16 and labels fit uint32.
             labels = relabel_sequential(labels.astype("uint32"))[0].astype("uint32")

micro_sam/v2/transforms/raw.py

@@ -392,10 +392,15 @@ def get_random_percentile_normalization(
     """
     augmentation2 = None
     if isinstance(raw_transform, RawTransform):
-        if not isinstance(raw_transform.normalizer, RandomPercentileNormalization):
-            raise ValueError(f"Unsupported generalist raw transform: {raw_transform!r}")
         augmentation1, augmentation2 = raw_transform.augmentation1, raw_transform.augmentation2
-        axis = raw_transform.normalizer.axis
+        if isinstance(raw_transform.normalizer, RandomPercentileNormalization):
+            axis = raw_transform.normalizer.axis
+        elif raw_transform.normalizer is _identity:
+            # Defect-augmented EM datasets (CREMI): the defect pipeline lives in augmentation1 with an
+            # identity normalizer. Preserve it and apply per-channel percentile normalization.
+            axis = (1, 2)
+        else:
+            raise ValueError(f"Unsupported normalizer in generalist raw transform: {raw_transform.normalizer!r}")
     elif isinstance(raw_transform, RandomPercentileNormalization):
         augmentation1, axis = None, raw_transform.axis
     elif raw_transform is _identity: