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:
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user