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")
# -- Audio ----------------------------------------------------------------
parser.add_argument("--sample_rate", type=int, default=24000, help="Audio sample rate (Hz)")
parser.add_argument("--num_mels", type=int, default=128, help="Number of mel bins (2.2: 128; up to 2.1: 64)")
parser.add_argument(
"--spec_width",
type=int,
default=None,
help=f"Frames per chunk (default per frontend: {', '.join(f'{k} {v}' for k, v in SPEC_WIDTH.items())})",
)
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=1280,
help="Embeddings layer size (2.2: 1280, N6-B0's top conv; up to 2.1: 1024)",
)
parser.add_argument("--alpha", type=float, default=1.5, help="Width multiplier")
parser.add_argument("--depth_multiplier", type=int, default=1, help="Depth multiplier")
parser.add_argument(
"--dw_kernel_size",
type=int,
default=5,
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(
"--backbone",
type=str,
default="efficientnet_b0_n6w",
choices=list(BACKBONES),
help=(
"Backbone behind the frontend: efficientnet_b0_n6w (2.3: N6-B0 with wider bottlenecks), efficientnet_b0_n6 "
"(2.2: EfficientNet-B0 with ReLU6, no squeeze-and-excite and a bounded residual stream) or dscnn "
"(releases up to 2.1; --alpha, --depth_multiplier, --dw_kernel_size and --stage_widths apply to it only)"
),
)
parser.add_argument(
"--raw_exposure_gain",
type=float,
default=None,
help=(
f"Gain of the raw filterbank's second, clamped exposure, mixed per band into a knee compressor "
f"(default {RAW_EXPOSURE_GAIN:g} for the raw frontend; 1 = one exposure)"
),
)
# -- 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(
"--crop_policy",
type=str,
default=None,
choices=["energy", "teacher"],
help=(
"How training chunks are chosen from a recording. 'energy' ranks candidates by "
"short-time energy and keeps the loudest; 'teacher' (needs --teacher_cache) keeps "
"the chunk where the teacher hears the labelled species most, and falls back to "
"energy where it has no score. Default: teacher with --teacher_cache, otherwise "
"energy. Evaluation is unaffected: it always scores whole files with overlapping windows."
),
)
parser.add_argument(
"--label_sidecar",
type=str,
default=None,
help=(
"CSV of per-file label additions for training (sample_id, positives, negatives; "
"classes separated by ';'): positives join the folder's class in the hard label "
"(multi-species recordings); negatives are confirmed absent, held at 0 and never "
"raised by the teacher's soft target. See birdnet_stm32.data.sidecar."
),
)
parser.add_argument(
"--teacher_cache",
type=str,
default=None,
help=(
"Directory of cached per-window teacher scores (see birdnet_stm32.data.teacher). "
"Each training chunk's label is blended with the teacher window nearest it. "
"Validation and checkpoint selection keep hard labels."
),
)
parser.add_argument(
"--teacher_weight",
type=float,
default=None,
help=(
f"Teacher share of the training target in [0, 1] (default {TEACHER_WEIGHT:g} with "
"--teacher_cache, otherwise 0 = hard labels only)"
),
)
parser.add_argument(
"--teacher_embedding_weight",
type=float,
default=None,
help=(
"Weight of a cosine loss between the student's embedding (through a training-only "
"projector) and the teacher's embedding of the same window. Default: "
f"{TEACHER_EMBEDDING_WEIGHT} when --teacher_cache holds emb.npy, otherwise off. "
"0 disables it. The deployed model is unchanged."
),
)
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(
"--class_cap_per_epoch",
type=int,
default=CLASS_CAP_PER_EPOCH,
help=(
f"At most this many training files per class (folder) in each epoch (default {CLASS_CAP_PER_EPOCH}; "
"0 = off). A larger "
"class contributes a different subset each epoch, rotating through all its files, so the draw is "
"more balanced without repeating minority files or discarding majority ones. With "
"--steps_per_epoch 0 an epoch is one such capped pass."
),
)
parser.add_argument(
"--warmup_epochs",
type=float,
default=TRAIN_WARMUP_EPOCHS,
help=(
f"Linear learning-rate warm-up before the cosine decay, in epochs (default {TRAIN_WARMUP_EPOCHS:g}; "
"may be fractional, e.g. 0.5 when one epoch is a pass over a large dataset)."
),
)
parser.add_argument(
"--steps_per_epoch",
type=int,
default=0,
help=(
"Training steps per epoch (default 0: one pass over the training files). Fixing it keeps a "
"schedule (warm-up, cosine decay, validation cadence) the same when the dataset grows."
),
)
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(
"--field_table",
type=str,
default=None,
help="CSV sample_id,site,recording of field clips: the class cap draws them for diversity (see --field_share)",
)
parser.add_argument(
"--field_share", type=float, default=0.5, help="Max share of a class's per-epoch cap filled by field clips"
)
parser.add_argument(
"--field_per_recording", type=int, default=2, help="Max segments of one field recording per class and epoch"
)
parser.add_argument("--resume", action="store_true", default=False, help="Resume training from checkpoint")
parser.add_argument(
"--init_checkpoint",
type=str,
default="",
help="Start from this model's weights at epoch 0; its frontend must match the architecture flags",
)
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
frontend = normalize_frontend_name(args.audio_frontend)
if args.spec_width is None:
args.spec_width = SPEC_WIDTH[frontend]
if args.raw_exposure_gain is None:
args.raw_exposure_gain = RAW_EXPOSURE_GAIN if frontend == "raw" else 1.0
if args.teacher_weight is None:
args.teacher_weight = TEACHER_WEIGHT if args.teacher_cache else 0.0
if args.crop_policy is None:
args.crop_policy = "teacher" if args.teacher_cache else "energy"
if not 0.0 <= args.teacher_weight <= 1.0:
parser.error(f"--teacher_weight must be in [0, 1], got {args.teacher_weight}")
if args.teacher_weight > 0 and not args.teacher_cache:
parser.error("--teacher_weight > 0 needs --teacher_cache")
has_embeddings = bool(args.teacher_cache) and os.path.isfile(os.path.join(args.teacher_cache, "emb.npy"))
if args.teacher_embedding_weight is None:
# On wherever the cache can feed it: float training with teacher embeddings.
float_training = not (args.qat or args.linear_probe)
args.teacher_embedding_weight = TEACHER_EMBEDDING_WEIGHT if has_embeddings and float_training else 0.0
if args.teacher_embedding_weight < 0:
parser.error(f"--teacher_embedding_weight must be >= 0, got {args.teacher_embedding_weight}")
if args.teacher_embedding_weight > 0 and not has_embeddings:
parser.error("--teacher_embedding_weight > 0 needs a --teacher_cache that holds emb.npy")
if args.field_table and args.class_cap_per_epoch <= 0:
parser.error("--field_table shapes the per-epoch class cap: it needs --class_cap_per_epoch > 0")
if args.teacher_embedding_weight > 0 and (args.qat or args.linear_probe):
parser.error("--teacher_embedding_weight applies to float training only (not --qat or --linear_probe)")
if args.crop_policy == "teacher" and not args.teacher_cache:
parser.error("--crop_policy teacher needs --teacher_cache")
args.deterministic = True # always deterministic
# Float training on a GPU runs XLA-compiled in mixed precision (~3x faster on
# the raw 448 model). QAT simulates the INT8 grid in float32, uncompiled.
float_training = not (args.qat or args.linear_probe)
has_gpu = bool(tf.config.list_physical_devices("GPU"))
args.mixed_precision = args.jit_compile = float_training and has_gpu
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