minimize log.out

Author: Donglai Wei

connectomics/training/lit/model.py

@@ -634,8 +634,15 @@ def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STEP_
                 print(f"[D1 Step {self.global_step}] TARGET: min={target_min:.3f}, max={target_max:.3f}, "
                       f"mean={target_mean:.3f}, >0: {target_positive_frac:.1f}%")
 
-        # Log losses (sync across GPUs for distributed training)
-        self.log_dict(loss_dict, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        # Keep full training curves in TensorBoard while avoiding console spam.
+        self.log_dict(
+            loss_dict,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=False,
+            logger=True,
+            sync_dist=False,
+        )
 
         return total_loss
 
@@ -728,8 +735,27 @@ def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STE
                         self.val_accuracy(preds, targets)
                         self.log('val_accuracy', self.val_accuracy, on_step=False, on_epoch=True, prog_bar=True)
 
-        # Log losses (sync across GPUs for distributed training)
-        self.log_dict(loss_dict, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        # Show only validation total loss on the progress bar.
+        if "val_loss_total" in loss_dict:
+            self.log(
+                "val_loss",
+                loss_dict["val_loss_total"],
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                logger=False,
+                sync_dist=True,
+            )
+
+        # Log full validation losses to logger at epoch granularity.
+        self.log_dict(
+            loss_dict,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=False,
+            logger=True,
+            sync_dist=True,
+        )
 
         return total_loss
 

connectomics/training/lit/trainer.py

@@ -19,7 +19,6 @@
     ModelCheckpoint,
     EarlyStopping,
     LearningRateMonitor,
-    RichProgressBar,
 )
 from pytorch_lightning.loggers import TensorBoardLogger
 from pytorch_lightning.plugins.environments import LightningEnvironment
@@ -106,7 +105,7 @@ def create_trainer(
             save_top_k=cfg.monitor.checkpoint.save_top_k,
             save_last=cfg.monitor.checkpoint.save_last,
             every_n_epochs=cfg.monitor.checkpoint.save_every_n_epochs,
-            verbose=True,
+            verbose=False,
             save_on_train_epoch_end=True,  # Save based on training metrics
         )
         callbacks.append(checkpoint_callback)
@@ -132,7 +131,7 @@ def create_trainer(
                 patience=cfg.monitor.early_stopping.patience,
                 mode=cfg.monitor.early_stopping.mode,
                 min_delta=cfg.monitor.early_stopping.min_delta,
-                verbose=True,
+                verbose=False,
                 check_on_train_epoch_end=True,  # Check at end of train epoch (not validation)
                 check_finite=cfg.monitor.early_stopping.check_finite,  # Stop on NaN/inf
                 stopping_threshold=cfg.monitor.early_stopping.threshold,
@@ -184,19 +183,13 @@ def create_trainer(
         # Previous fix in val_dataloader() only ran once during setup
         validation_reseeding_callback = ValidationReseedingCallback(
             base_seed=cfg.system.seed,
-            log_fingerprint=True,
+            log_fingerprint=False,
             log_all_ranks=False,
-            verbose=True,
+            verbose=False,
         )
         callbacks.append(validation_reseeding_callback)
         print(f"  Validation Reseeding: Enabled (base_seed={cfg.system.seed})")
 
-    # Progress bar (optional - requires rich package)
-    try:
-        callbacks.append(RichProgressBar())
-    except (ImportError, ModuleNotFoundError):
-        pass  # Use default progress bar
-
     # Setup logger (training only - in run_dir/logs/)
     # Always create a logger for training to avoid warnings about missing logger
     logger = None
@@ -322,6 +315,7 @@ def create_trainer(
         benchmark=cfg.optimization.benchmark,
         fast_dev_run=bool(fast_dev_run),
         detect_anomaly=detect_anomaly,
+        enable_progress_bar=False,
         plugins=plugins,
     )
 

scripts/main.py

@@ -26,6 +26,7 @@
     python scripts/main.py --config tutorials/mito_lucchi++.yaml --checkpoint path/to/ckpt.ckpt --reset-max-epochs 500
 """
 
+import os
 import sys
 from pathlib import Path
 
@@ -68,6 +69,23 @@
 # Setup seed_everything with version fallback
 seed_everything = setup_seed_everything()
 
+_RANK_STDOUT_REDIRECT = None
+
+
+def suppress_nonzero_rank_stdout() -> None:
+    """Reduce duplicate stdout spam from DDP subprocesses.
+
+    In local multi-GPU spawn, each subprocess executes this script and prints
+    the same setup logs. Keep rank 0 stdout visible and silence stdout on
+    non-zero ranks. stderr is untouched for error visibility.
+    """
+    global _RANK_STDOUT_REDIRECT
+    local_rank = os.environ.get("LOCAL_RANK")
+    if local_rank is None or local_rank == "0":
+        return
+    _RANK_STDOUT_REDIRECT = open(os.devnull, "w")
+    sys.stdout = _RANK_STDOUT_REDIRECT
+
 
 def configure_matmul_precision(cfg: Config) -> None:
     """Enable Tensor Core matmul precision when supported by available CUDA devices."""
@@ -164,6 +182,8 @@ def extract_step_from_checkpoint(checkpoint_path: str) -> str:
 
 def main():
     """Main training function."""
+    suppress_nonzero_rank_stdout()
+
     # Parse arguments
     args = parse_args()