def get_args() -> argparse.Namespace:
"""Parse command-line arguments for training.
Sensible defaults are chosen so that most users only need to specify
``--data_path_train``. SpecAugment, deterministic seeding, and gradient
clipping are enabled by default.
"""
parser = argparse.ArgumentParser(description="Train STM32N6 audio classifier")
# -- Data -----------------------------------------------------------------
parser.add_argument("--data_path_train", type=str, required=True, help="Path to train dataset")
parser.add_argument(
"--data_path_val",
type=str,
default=None,
help="Separate validation dataset root; disables the random --val_split",
)
parser.add_argument(
"--classes_file",
type=str,
default=None,
help="Ordered one-label-per-line output schema; noise remains an all-zero folder",
)
parser.add_argument(
"--model_config",
type=str,
default="",
help=(
"Architecture config for --qat and --linear_probe. Defaults to the "
"checkpoint's sibling _model_config.json; pass it explicitly when fine-tuning a "
"checkpoint that inherited its architecture from an earlier step, such as a QAT "
"checkpoint, which writes no config of its own."
),
)
parser.add_argument("--max_classes", type=int, default=None, help="Use top N classes by sample count")
parser.add_argument("--max_samples", type=int, default=None, help="Max samples per class")
parser.add_argument("--upsample_ratio", type=float, default=0.5, help="Upsample ratio for minority classes")
# -- Audio ----------------------------------------------------------------
parser.add_argument("--sample_rate", type=int, default=24000, help="Audio sample rate (Hz)")
parser.add_argument("--num_mels", type=int, default=64, help="Number of mel bins")
parser.add_argument("--spec_width", type=int, default=256, help="Spectrogram width (frames)")
parser.add_argument("--fft_length", type=int, default=512, help="FFT length")
parser.add_argument("--chunk_duration", type=float, default=2.5, help="Audio chunk duration (seconds)")
parser.add_argument(
"--max_duration",
type=int,
default=60,
help=(
"Maximum seconds to read per file. The loader still reads only the bytes it needs "
"for the candidate chunks (smart-crop bounded by --max_chunks_per_file)."
),
)
parser.add_argument(
"--audio_frontend",
type=str,
default="raw",
choices=["raw", "hybrid", "librosa"],
help="Audio frontend mode",
)
parser.add_argument(
"--mag_scale",
type=str,
default="pwl",
choices=["pwl", "none"],
help="Magnitude scaling: learned hinge sum (pwl) or none",
)
parser.add_argument(
"--input_compression",
type=str,
default="none",
choices=["none", "sqrt", "log"],
help="Compress a librosa/hybrid spectrogram input before it is quantized (computed on host/M55)",
)
# -- Model architecture ---------------------------------------------------
parser.add_argument("--embeddings_size", type=int, default=512, help="Embeddings layer size")
parser.add_argument("--alpha", type=float, default=1.0, help="Width multiplier")
parser.add_argument("--depth_multiplier", type=int, default=1, help="Depth multiplier")
parser.add_argument(
"--head_pooling",
type=str,
default="gap",
choices=list(HEAD_POOLINGS),
help="Pooling head: global average (gap), or mean over frequency then max + mean over time",
)
parser.add_argument(
"--dw_kernel_size",
type=int,
default=3,
choices=list(DW_KERNEL_SIZES),
help="Depthwise kernel size in stages 2-4 (stage 1 stays 3x3)",
)
parser.add_argument(
"--stage_widths",
type=int,
nargs=4,
default=list(STAGE_WIDTHS),
metavar="W",
help="Base output channels of the four stages, before --alpha",
)
parser.add_argument("--frontend_trainable", action="store_true", default=False)
# -- Augmentation ---------------------------------------------------------
parser.add_argument("--no_spec_augment", action="store_true", default=False, help="Disable SpecAugment")
parser.add_argument("--freq_mask_max", type=int, default=8, help="Max frequency mask width (bins)")
parser.add_argument("--time_mask_max", type=int, default=25, help="Max time mask width (frames)")
parser.add_argument("--mixup_alpha", type=float, default=0.2, help="Mixup alpha")
parser.add_argument("--mixup_probability", type=float, default=0.25, help="Mixup batch fraction")
# -- Training -------------------------------------------------------------
parser.add_argument("--batch_size", type=int, default=32, help="Batch size")
parser.add_argument("--num_workers", type=int, default=8, help="Parallel data loading workers (0 = sequential)")
parser.add_argument(
"--max_chunks_per_file",
type=int,
default=1,
help="Max salient chunks to extract per file open (reduces redundant I/O for long recordings)",
)
parser.add_argument(
"--prefetch_batches",
type=int,
default=2,
help="Loader prefetch queue depth in batches (higher = faster, but more RAM)",
)
parser.add_argument(
"--epochs",
type=int,
default=None,
help=f"Number of epochs (default: {TRAIN_EPOCHS}; {QAT_EPOCHS} with --qat)",
)
parser.add_argument(
"--learning_rate",
type=float,
default=None,
help=(
f"Initial learning rate (default: {TRAIN_LEARNING_RATE:g}; {QAT_LEARNING_RATE:g} with --qat, "
f"{PROBE_LEARNING_RATE:g} with --linear_probe)"
),
)
parser.add_argument("--dropout", type=float, default=0.5, help="Dropout rate before classifier head")
parser.add_argument("--optimizer", type=str, default="adam", choices=["adam", "sgd", "adamw"], help="Optimizer")
parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay (adamw only)")
parser.add_argument("--val_split", type=float, default=0.2, help="Validation split ratio")
parser.add_argument(
"--checkpoint_path", type=str, default="checkpoints/best_model.keras", help="Output checkpoint path (.keras)"
)
parser.add_argument("--grad_clip", type=float, default=1.0, help="Max gradient norm for clipping (0 = disabled)")
parser.add_argument(
"--mixed_precision", action="store_true", default=False, help="Enable FP16 mixed precision training"
)
parser.add_argument("--resume", action="store_true", default=False, help="Resume training from checkpoint")
parser.add_argument("--seed", type=int, default=42, help="Random seed for deterministic training")
parser.add_argument(
"--validation_overlap",
type=float,
default=None,
help="File validation overlap in seconds (default: half the model chunk duration).",
)
parser.add_argument(
"--validation_pooling",
choices=["max", "avg", "lme"],
default="max",
help="File validation pooling for exact cMAP checkpoint selection.",
)
parser.add_argument(
"--validation_subset",
type=int,
default=0,
help=(
"Score checkpoint selection (standard training and QAT) on a fixed stratified subset of "
"this many validation files instead of all of them. Classes are drawn round-robin with "
"seed 1234, so the subset is as class-balanced as the folders allow and identical across "
"runs and epochs. Subset cMAP is biased upward and not comparable to full-manifest "
"numbers. 0 uses all."
),
)
# -- QAT ------------------------------------------------------------------
parser.add_argument(
"--qat",
action="store_true",
default=False,
help="Quantization-aware fine-tuning (requires pretrained --checkpoint_path)",
)
parser.add_argument(
"--qat_calibration_samples",
type=int,
default=1024,
help="Exact stratified samples used for QAT ranges and final INT8 calibration",
)
parser.add_argument(
"--qat_distillation_weight",
type=float,
default=1.0,
help="QAT teacher Bernoulli-KL loss weight",
)
parser.add_argument(
"--qat_cosine_weight",
type=float,
default=0.10,
help="QAT mean teacher/student cosine-loss weight",
)
parser.add_argument(
"--qat_cosine_tail_weight",
type=float,
default=0.75,
help="QAT worst-sample teacher/student cosine-loss weight",
)
parser.add_argument(
"--qat_cosine_tail_fraction",
type=float,
default=0.10,
help="Fraction of each QAT batch included in the worst-sample loss",
)
parser.add_argument(
"--qat_range_refresh",
action=argparse.BooleanOptionalAction,
default=True,
help="Recalibrate QAT activation ranges on the current weights after every epoch",
)
# -- Linear probing -------------------------------------------------------
parser.add_argument(
"--linear_probe",
action="store_true",
default=False,
help="Freeze backbone and train only the classifier head (requires pretrained --checkpoint_path)",
)
args = parser.parse_args()
# Derive positive flags from --no_* flags
args.spec_augment = not args.no_spec_augment
args.deterministic = True # always deterministic
if args.validation_subset < 0:
parser.error("--validation_subset must be non-negative")
# The compression steps run one at a time against a converged checkpoint;
# the documented order is QAT, then convert.
exclusive = [name for name in ("qat", "linear_probe") if getattr(args, name)]
if len(exclusive) > 1:
raise SystemExit(f"Options are mutually exclusive, run them as separate steps: {exclusive}")
# Each step has its own best schedule: fine-tuning at the training rate
# destroys a converged model (measured for QAT at 2e-4).
if args.epochs is None:
args.epochs = QAT_EPOCHS if args.qat else TRAIN_EPOCHS
if args.learning_rate is None:
args.learning_rate = (
QAT_LEARNING_RATE if args.qat else PROBE_LEARNING_RATE if args.linear_probe else TRAIN_LEARNING_RATE
)
return args