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>
This commit is contained in:
2026-06-21 20:32:01 +08:00
parent 02d429282e
commit b0942180ba
+12 -6
View File
@@ -114,6 +114,8 @@ def main():
help="Path to checkpoint .pt file to resume training from") help="Path to checkpoint .pt file to resume training from")
parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True, parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True,
help="Enable Automatic Mixed Precision (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() args = parser.parse_args()
use_amp = args.amp 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") device = torch.device(args.device if torch.cuda.is_available() and "cuda" in args.device else "cpu")
print(f"Device: {device}") 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 # Create model
model = VelocityPredictionModel() model = VelocityPredictionModel()
model.to(device) model.to(device)
@@ -135,7 +142,7 @@ def main():
stride=train_cfg.sliding_window_stride, stride=train_cfg.sliding_window_stride,
batch_size=train_cfg.batch_size, batch_size=train_cfg.batch_size,
num_workers=train_cfg.num_workers, num_workers=train_cfg.num_workers,
event_threshold=train_cfg.event_threshold, event_threshold=event_threshold,
event_use_log=train_cfg.event_use_log, event_use_log=train_cfg.event_use_log,
) )
val_loader = create_val_loader( val_loader = create_val_loader(
@@ -143,7 +150,7 @@ def main():
stride=train_cfg.sliding_window_stride, stride=train_cfg.sliding_window_stride,
batch_size=train_cfg.batch_size, batch_size=train_cfg.batch_size,
num_workers=train_cfg.num_workers, num_workers=train_cfg.num_workers,
event_threshold=train_cfg.event_threshold, event_threshold=event_threshold,
event_use_log=train_cfg.event_use_log, event_use_log=train_cfg.event_use_log,
) )
@@ -161,9 +168,6 @@ def main():
# criterion = nn.SmoothL1Loss() # criterion = nn.SmoothL1Loss()
criterion = nn.MSELoss() criterion = nn.MSELoss()
ckpt_dir = Path(train_cfg.checkpoint_dir)
ckpt_dir.mkdir(parents=True, exist_ok=True)
# ── Resume from checkpoint ──────────────────────────────────── # ── Resume from checkpoint ────────────────────────────────────
start_epoch = 1 start_epoch = 1
global_step = 0 global_step = 0
@@ -191,11 +195,13 @@ def main():
else: else:
print(f"\nStarting training for {train_cfg.epochs} epochs...") 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: if run_id is None:
run_id = time.strftime("run_%Y%m%d_%H%M%S") run_id = time.strftime("run_%Y%m%d_%H%M%S")
log_dir = Path(train_cfg.log_dir) / run_id log_dir = Path(train_cfg.log_dir) / run_id
log_dir.mkdir(parents=True, exist_ok=True) 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)) writer = SummaryWriter(log_dir=str(log_dir))
print(f" seq_len={train_cfg.seq_len}, batch_size={train_cfg.batch_size}") print(f" seq_len={train_cfg.seq_len}, batch_size={train_cfg.batch_size}")