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