| """Deterministic train/validation-only training for the DeepWeeds comparison.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import random |
| import sys |
| import time |
| from collections.abc import Mapping, Sequence |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| from safetensors.torch import save_file |
| from torch import nn |
| from torch.optim import Adam, Optimizer |
| from torch.utils.data import DataLoader |
|
|
| try: |
| from .data import build_comparison_transforms, load_fold_splits, make_loader |
| from .evaluate import compute_classification_metrics, package_environment, resolve_device, write_json |
| from .model import SmallDeepWeedsCNN |
| except ImportError: |
| from data import build_comparison_transforms, load_fold_splits, make_loader |
| from evaluate import compute_classification_metrics, package_environment, resolve_device, write_json |
| from model import SmallDeepWeedsCNN |
|
|
|
|
| @dataclass(frozen=True) |
| class TrainingConfig: |
| """The immutable optimization contract shared by good and bad runs.""" |
|
|
| seed: int = 2026 |
| batch_size: int = 64 |
| epochs: int = 12 |
| learning_rate: float = 1e-3 |
| optimizer: str = "Adam" |
| scheduler: None = None |
| loss: str = "CrossEntropyLoss" |
| class_weights: None = None |
| checkpoint_selection: str = "minimum_validation_loss" |
| fold: int = 0 |
|
|
|
|
| def training_config_for_run(run_kind: str) -> TrainingConfig: |
| """Return the same declared training configuration for either arm.""" |
|
|
| if run_kind not in {"good", "bad"}: |
| raise ValueError("run_kind must be 'good' or 'bad'") |
| return TrainingConfig() |
|
|
|
|
| def build_criterion() -> nn.CrossEntropyLoss: |
| """Build ordinary, unweighted multiclass cross-entropy.""" |
|
|
| return nn.CrossEntropyLoss(weight=None) |
|
|
|
|
| def build_optimizer(model: nn.Module, config: TrainingConfig) -> Adam: |
| """Build the sole optimizer declared by the comparison contract.""" |
|
|
| if config.optimizer != "Adam" or config.learning_rate != 1e-3: |
| raise ValueError("The comparison requires Adam with learning rate 1e-3") |
| return Adam(model.parameters(), lr=config.learning_rate) |
|
|
|
|
| def seed_everything(seed: int) -> None: |
| """Seed supported random sources and deterministic CUDA behavior.""" |
|
|
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.benchmark = False |
| torch.backends.cudnn.deterministic = True |
| torch.use_deterministic_algorithms(True) |
|
|
|
|
| def select_checkpoint_by_validation_loss( |
| validation_losses: Mapping[str, float], |
| ) -> str: |
| """Return the key with the lowest finite validation loss and no other input.""" |
|
|
| if not validation_losses: |
| raise ValueError("At least one validation loss is required") |
| if any(not math.isfinite(float(loss)) for loss in validation_losses.values()): |
| raise ValueError("Validation losses must be finite") |
| return min(validation_losses, key=validation_losses.__getitem__) |
|
|
|
|
| def _run_epoch( |
| model: nn.Module, |
| loader: DataLoader, |
| criterion: nn.CrossEntropyLoss, |
| device: torch.device, |
| optimizer: Optimizer | None, |
| ) -> dict[str, float]: |
| training = optimizer is not None |
| model.train(training) |
| loss_sum = 0.0 |
| sample_count = 0 |
| true_labels: list[int] = [] |
| predicted_labels: list[int] = [] |
| class_count: int | None = None |
|
|
| with torch.set_grad_enabled(training): |
| for images, labels, _filenames in loader: |
| images = images.to(device, non_blocking=device.type == "cuda") |
| labels = labels.to(device, non_blocking=device.type == "cuda") |
| if optimizer is not None: |
| optimizer.zero_grad(set_to_none=True) |
| logits = model(images) |
| if logits.ndim != 2: |
| raise ValueError(f"Expected two-dimensional logits, got {tuple(logits.shape)}") |
| class_count = logits.shape[1] |
| loss = criterion(logits, labels) |
| if not torch.isfinite(loss): |
| raise FloatingPointError(f"Encountered non-finite {'train' if training else 'validation'} loss") |
| if optimizer is not None: |
| loss.backward() |
| optimizer.step() |
|
|
| batch_size = labels.shape[0] |
| sample_count += batch_size |
| loss_sum += float(loss.detach().item()) * batch_size |
| true_labels.extend(int(label) for label in labels.detach().cpu().tolist()) |
| predicted_labels.extend(int(label) for label in logits.detach().argmax(dim=1).cpu().tolist()) |
|
|
| if sample_count == 0 or class_count is None: |
| raise ValueError("Training and validation loaders must not be empty") |
| metrics = compute_classification_metrics( |
| true_labels, |
| predicted_labels, |
| class_names=[str(class_id) for class_id in range(class_count)], |
| ) |
| return { |
| "loss": loss_sum / sample_count, |
| "accuracy": metrics["accuracy"], |
| "macro_f1": metrics["macro_f1"], |
| "balanced_accuracy": metrics["balanced_accuracy"], |
| } |
|
|
|
|
| def _save_safetensors(model: nn.Module, path: Path) -> None: |
| state = { |
| name: tensor.detach().cpu().contiguous() |
| for name, tensor in model.state_dict().items() |
| } |
| save_file(state, str(path)) |
|
|
|
|
| def train_model( |
| model: nn.Module, |
| train_loader: DataLoader, |
| validation_loader: DataLoader, |
| output_dir: str | Path, |
| device: torch.device, |
| config: TrainingConfig, |
| ) -> dict[str, Any]: |
| """Fit from train data and select from validation loss; no test input exists.""" |
|
|
| if config != TrainingConfig(): |
| raise ValueError("TrainingConfig must match the fixed comparison contract") |
| destination = Path(output_dir).resolve() |
| destination.mkdir(parents=True, exist_ok=True) |
| checkpoint_path = destination / "model.safetensors" |
| seed_everything(config.seed) |
| model.to(device) |
| criterion = build_criterion() |
| optimizer = build_optimizer(model, config) |
| history: list[dict[str, Any]] = [] |
| validation_losses: dict[str, float] = {} |
| started = time.perf_counter() |
|
|
| for epoch in range(1, config.epochs + 1): |
| train_metrics = _run_epoch(model, train_loader, criterion, device, optimizer) |
| validation_metrics = _run_epoch(model, validation_loader, criterion, device, None) |
| history.append( |
| { |
| "epoch": epoch, |
| "train": train_metrics, |
| "validation": validation_metrics, |
| } |
| ) |
| epoch_key = f"epoch_{epoch}" |
| validation_losses[epoch_key] = validation_metrics["loss"] |
| if select_checkpoint_by_validation_loss(validation_losses) == epoch_key: |
| _save_safetensors(model, checkpoint_path) |
|
|
| elapsed = time.perf_counter() - started |
| best_key = select_checkpoint_by_validation_loss(validation_losses) |
| best_epoch = int(best_key.removeprefix("epoch_")) |
| result = { |
| "history": history, |
| "best_epoch": best_epoch, |
| "best_validation_loss": validation_losses[best_key], |
| "checkpoint_path": str(checkpoint_path), |
| "checkpoint_selection": config.checkpoint_selection, |
| "runtime": {"training_seconds": elapsed}, |
| "environment": package_environment(device), |
| } |
| write_json(destination / "history.json", result) |
| return result |
|
|
|
|
| def run_training( |
| run_kind: str, |
| labels_dir: str | Path, |
| images_dir: str | Path, |
| output_dir: str | Path, |
| device: str = "auto", |
| num_workers: int = 4, |
| ) -> dict[str, Any]: |
| """Create official fold-0 train/validation loaders and train one declared arm.""" |
|
|
| config = training_config_for_run(run_kind) |
| if num_workers < 0: |
| raise ValueError("num_workers must be non-negative") |
| labels_path = Path(labels_dir).resolve() |
| images_path = Path(images_dir).resolve() |
| output_path = Path(output_dir).resolve() |
| selected_device = resolve_device(device) |
| splits = load_fold_splits(labels_dir=labels_path, images_dir=images_path, fold=config.fold) |
| transforms = build_comparison_transforms() |
| pin_memory = selected_device.type == "cuda" |
| train_loader = make_loader( |
| splits["train"], |
| transforms[f"{run_kind}_train"], |
| shuffle=True, |
| seed=config.seed, |
| batch_size=config.batch_size, |
| num_workers=num_workers, |
| pin_memory=pin_memory, |
| ) |
| validation_loader = make_loader( |
| splits["val"], |
| transforms[f"{run_kind}_val"], |
| shuffle=False, |
| seed=config.seed, |
| batch_size=config.batch_size, |
| num_workers=num_workers, |
| pin_memory=pin_memory, |
| ) |
| seed_everything(config.seed) |
| model = SmallDeepWeedsCNN( |
| num_classes=9, |
| in_channels=3 if run_kind == "good" else 1, |
| ) |
| training_args = { |
| "training": asdict(config), |
| "run": { |
| "kind": run_kind, |
| "device": str(selected_device), |
| "num_workers": num_workers, |
| "pin_memory": pin_memory, |
| "fold_files": { |
| "train": str((labels_path / "train_subset0.csv").resolve()), |
| "validation": str((labels_path / "val_subset0.csv").resolve()), |
| }, |
| "images_dir": str(images_path), |
| }, |
| } |
| write_json(output_path / "training_args.json", training_args) |
| result = train_model( |
| model, |
| train_loader, |
| validation_loader, |
| output_path, |
| selected_device, |
| config, |
| ) |
| result["run"] = training_args["run"] |
| return result |
|
|
|
|
| def _parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--run-kind", choices=("good", "bad"), required=True) |
| parser.add_argument("--labels-dir", type=Path, required=True) |
| parser.add_argument("--images-dir", type=Path, required=True) |
| parser.add_argument("--output-dir", type=Path, required=True) |
| parser.add_argument("--device", default="auto") |
| parser.add_argument("--num-workers", type=int, default=4) |
| return parser |
|
|
|
|
| def main(argv: Sequence[str] | None = None) -> int: |
| args = _parser().parse_args(argv) |
| result = run_training( |
| run_kind=args.run_kind, |
| labels_dir=args.labels_dir, |
| images_dir=args.images_dir, |
| output_dir=args.output_dir, |
| device=args.device, |
| num_workers=args.num_workers, |
| ) |
| print( |
| json.dumps( |
| { |
| "run_kind": args.run_kind, |
| "best_epoch": result["best_epoch"], |
| "best_validation_loss": result["best_validation_loss"], |
| "checkpoint_path": result["checkpoint_path"], |
| "training_seconds": result["runtime"]["training_seconds"], |
| }, |
| indent=2, |
| allow_nan=False, |
| ) |
| ) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
