From b0942180bad026337d1f9f82f6e7c9b092007c3f Mon Sep 17 00:00:00 2001
From: CaoWangrenbo <cao.wangrenbo@yandex.com>
Date: Sun, 21 Jun 2026 20:32:01 +0800
Subject: [PATCH] feat: add --threshold CLI arg for event brightness override;
 move ckpt_dir under run_id subdirectory

- --threshold CLI arg overrides config event_threshold at runtime
- Move ckpt_dir creation after run_id resolution (run-specific subdirectory)
- Use overridden threshold in train/val loader creation

Generated by Mistral Vibe.
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
---
 src/velocity_prediction/train.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/velocity_prediction/train.py b/src/velocity_prediction/train.py
index 62b6b93..8e22fef 100644
--- a/src/velocity_prediction/train.py
+++ b/src/velocity_prediction/train.py
@@ -114,6 +114,8 @@ def main():
                         help="Path to checkpoint .pt file to resume training from")
     parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True,
                         help="Enable Automatic Mixed Precision (default: True)")
+    parser.add_argument("--threshold", type=float, default=None,
+                        help="Event brightness change threshold (overrides config, default: same as train_cfg)")
     args = parser.parse_args()
     use_amp = args.amp
 
@@ -121,6 +123,11 @@ def main():
     device = torch.device(args.device if torch.cuda.is_available() and "cuda" in args.device else "cpu")
     print(f"Device: {device}")
 
+    # Override threshold from CLI if provided
+    event_threshold = args.threshold if args.threshold is not None else train_cfg.event_threshold
+    if args.threshold is not None:
+        print(f"Event threshold overridden: {train_cfg.event_threshold} → {event_threshold}")
+
     # Create model
     model = VelocityPredictionModel()
     model.to(device)
@@ -135,7 +142,7 @@ def main():
         stride=train_cfg.sliding_window_stride,
         batch_size=train_cfg.batch_size,
         num_workers=train_cfg.num_workers,
-        event_threshold=train_cfg.event_threshold,
+        event_threshold=event_threshold,
         event_use_log=train_cfg.event_use_log,
     )
     val_loader = create_val_loader(
@@ -143,7 +150,7 @@ def main():
         stride=train_cfg.sliding_window_stride,
         batch_size=train_cfg.batch_size,
         num_workers=train_cfg.num_workers,
-        event_threshold=train_cfg.event_threshold,
+        event_threshold=event_threshold,
         event_use_log=train_cfg.event_use_log,
     )
 
@@ -161,9 +168,6 @@ def main():
     # criterion = nn.SmoothL1Loss()
     criterion = nn.MSELoss()
 
-    ckpt_dir = Path(train_cfg.checkpoint_dir)
-    ckpt_dir.mkdir(parents=True, exist_ok=True)
-
     # ── Resume from checkpoint ────────────────────────────────────
     start_epoch = 1
     global_step = 0
@@ -191,11 +195,13 @@ def main():
     else:
         print(f"\nStarting training for {train_cfg.epochs} epochs...")
 
-    # Logging — run-specific subdirectory for isolation + resume continuity
+    # Logging & checkpoint — run-specific subdirectories for isolation + resume continuity
     if run_id is None:
         run_id = time.strftime("run_%Y%m%d_%H%M%S")
     log_dir = Path(train_cfg.log_dir) / run_id
     log_dir.mkdir(parents=True, exist_ok=True)
+    ckpt_dir = Path(train_cfg.checkpoint_dir) / run_id
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
     writer = SummaryWriter(log_dir=str(log_dir))
 
     print(f"  seq_len={train_cfg.seq_len}, batch_size={train_cfg.batch_size}")