Buckets:
| """Bounded source-aware semantic fit interface, without learning-gate orchestration.""" | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from copy import deepcopy | |
| import json | |
| import math | |
| import os | |
| import random | |
| import time | |
| import torch | |
| from load_dataset import dataset, load_sample | |
| from training.model_factory import ROOT, load_pretrained_segmenter | |
| from training.runtime import load_runtime_config | |
| from training.segmentation_data import (CANONICAL_IDS, convert_segmentation_sample, | |
| transform_example, segmentation_batch) | |
| from training.detection_trainer import freeze_batchnorm_statistics, batchnorm_buffers | |
| from training.segmentation_loss import (SourceAwareSemanticLoss, build_adapted_segmenter, | |
| segmentation_counts, summarize_counts) | |
| class SegmenterTrainConfig: | |
| root: Path = ROOT / "datasets/corpus" | |
| output: Path = ROOT / "training/runs/segmenter" | |
| learning_rate: float = .0001 | |
| epochs: int = 10 | |
| patience: int = 3 | |
| seed: int = 42 | |
| record_initial_metrics: bool = False | |
| expected_initial_sha256: str | None = None | |
| freeze_bn_statistics: bool = False | |
| def from_json(cls, path, **overrides): | |
| data = json.loads(path.read_text()) | |
| fixed = {"checkpoint": "yolo26n-sem-ade20k.pt", "image_size": 1024, "batch_size": 1, | |
| "device": "cpu", "use_amp": False, "canonical_ids": list(CANONICAL_IDS), | |
| "license_scope": "all", "optimizer": "AdamW", | |
| "augmentation": {"horizontal_flip": True, "vertical_flip": True, "right_angle_rotations": True}} | |
| if set(data) != set(fixed) | {"learning_rate", "epochs", "patience", "seed"}: | |
| raise ValueError("unexpected semantic configuration fields") | |
| if any(data[k] != value or type(data[k]) is not type(value) for k, value in fixed.items()): | |
| raise ValueError("semantic config changes the approved policy") | |
| return cls(**{**{k: data[k] for k in ("learning_rate", "epochs", "patience", "seed")}, **overrides}) | |
| def __post_init__(self): | |
| if type(self.freeze_bn_statistics) is not bool: | |
| raise ValueError("freeze_bn_statistics must be boolean") | |
| if not Path(self.output).resolve().is_relative_to(ROOT / "training"): | |
| raise ValueError("semantic outputs must remain under training/") | |
| if isinstance(self.learning_rate, bool) or not math.isfinite(self.learning_rate) or self.learning_rate <= 0: | |
| raise ValueError("positive finite learning rate required") | |
| for value, low, high in [(self.epochs, 1, 10), (self.patience, 1, 3), (self.seed, 0, 2**32 - 1)]: | |
| if type(value) is not int or not low <= value <= high: | |
| raise ValueError("invalid epoch/patience/seed setting") | |
| class Trainer: | |
| def __init__(self, config): | |
| self.config = config | |
| def fit(self, train_ids: list[str], validation_ids: list[str], *, max_seconds: int) -> dict: | |
| return self._fit(train_ids, validation_ids, max_seconds=max_seconds, diagnostic=False) | |
| def fit_overfit(self, sample_ids: list[str], *, max_seconds: int) -> dict: | |
| if len(sample_ids) != 24 or len(set(sample_ids)) != 24 or not 0 < max_seconds <= 900: | |
| raise ValueError("overfit requires exactly24 unique train tiles and <=900 seconds") | |
| return self._fit(sample_ids, sample_ids, max_seconds=max_seconds, diagnostic=True) | |
| def _fit(self, train_ids, validation_ids, *, max_seconds, diagnostic): | |
| if type(max_seconds) is not int or not 0 < max_seconds <= 3600: | |
| raise ValueError("fit budget must be in (0, 3600]") | |
| config = self.config | |
| train = dataset(config.root, "segmentation", "train", license_scope="all") | |
| val = train if diagnostic else dataset(config.root, "segmentation", "val", license_scope="all") | |
| training, validation = {r["id"]: r for r in train.samples}, {r["id"]: r for r in val.samples} | |
| manifest = json.loads((ROOT / "evaluation/quick_eval_manifest.json").read_text()) | |
| reserved = {r["sample_id"] for r in manifest["selection_validation"] if r["task"] == "segmentation"} | |
| for ids, allowed in [(train_ids, set(training)), (validation_ids, set(training) if diagnostic else set(validation) & reserved)]: | |
| if not ids or len(ids) != len(set(ids)) or not set(ids) <= allowed: | |
| raise ValueError("unique train IDs and reserved selection-validation IDs required") | |
| output = Path(config.output).resolve() | |
| output.mkdir(parents=True, exist_ok=False) | |
| started = time.monotonic() | |
| def budget(): | |
| if time.monotonic() - started >= max_seconds: | |
| raise TimeoutError("cooperative semantic fit budget exhausted") | |
| runtime_path = Path(os.environ.get("LONGMEMORY_GEO_RUNTIME_CONFIG", ROOT / "training/configs/runtime.json")) | |
| runtime = load_runtime_config(runtime_path) | |
| device = runtime.device | |
| if device == "cpu": | |
| torch.set_num_threads(runtime.num_threads) | |
| torch.manual_seed(config.seed) | |
| if device.startswith("cuda"): | |
| torch.cuda.manual_seed_all(config.seed) | |
| rng = random.Random(config.seed) | |
| history, best_metrics, best_checkpoint = [], {}, None | |
| status, stale, best = "COMPLETED", 0, -1. | |
| try: | |
| budget() | |
| pretrained = load_pretrained_segmenter(ROOT / "training/cache/checkpoints", device=device) | |
| model, transfer = build_adapted_segmenter(pretrained, seed=config.seed) | |
| model.to(device) | |
| model.criterion = SourceAwareSemanticLoss(model) | |
| for key, value in {"imgsz": 1024, "batch": 1, "device": "cpu", "amp": False, | |
| "seed": config.seed, "epochs": config.epochs, "patience": config.patience, | |
| "optimizer": "AdamW", "lr0": config.learning_rate, "weight_decay": .01}.items(): | |
| setattr(model.args, key, device if key == "device" else (runtime.use_amp if key == "amp" else value)) | |
| optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate) | |
| predictor = deepcopy(pretrained) | |
| from evaluation.quick_eval import semantic_array | |
| from training.model_factory import model_state_sha256 | |
| initial_sha = model_state_sha256(model) | |
| initial_bn = batchnorm_buffers(model) if config.freeze_bn_statistics else {} | |
| if config.expected_initial_sha256 is not None and initial_sha != config.expected_initial_sha256: | |
| raise ValueError("adapted initialization differs from approved baseline") | |
| def evaluate(): | |
| model.eval() | |
| predictor.model, predictor.predictor = deepcopy(model), None | |
| counts = {str(c): {"intersection": 0, "union": 0, "targets": 0} for c in CANONICAL_IDS} | |
| for sample_id in validation_ids: | |
| budget() | |
| sample = convert_segmentation_sample(load_sample(config.root, validation[sample_id], val.expected_hashes)) | |
| with torch.inference_mode(): | |
| result = predictor.predict(sample.image[:, :, ::-1].copy(), imgsz=1024, | |
| device=device, save=False, verbose=False)[0] | |
| from training.segmentation_data import head_to_canonical | |
| prediction = head_to_canonical(semantic_array(result)) | |
| current = segmentation_counts(prediction, sample.mask) | |
| for c in counts: | |
| for field in counts[c]: | |
| counts[c][field] += current[c][field] | |
| budget() | |
| metrics = summarize_counts(counts) | |
| if metrics["miou"] is None: | |
| raise ValueError("selection validation contains no labeled pixels") | |
| return metrics | |
| if config.record_initial_metrics or diagnostic: | |
| history.append({"epoch": 0, "metrics": evaluate(), "initial_state_sha256": initial_sha, | |
| "evaluation_ids": validation_ids, "diagnostic_train_only": diagnostic, | |
| "elapsed_seconds": time.monotonic() - started}) | |
| (output / "initial_evaluation.json").write_text(json.dumps(history[0], indent=2), encoding="utf-8") | |
| initial_progress = {"status": "IN_PROGRESS", "best_checkpoint": None, "history": history, | |
| "best_metrics": {}, "elapsed_seconds": time.monotonic() - started} | |
| temporary_progress = output / "progress.json.partial" | |
| temporary_progress.write_text(json.dumps(initial_progress, indent=2), encoding="utf-8") | |
| temporary_progress.replace(output / "progress.json") | |
| for epoch in range(config.epochs): | |
| budget() | |
| model.train() | |
| if config.freeze_bn_statistics: | |
| freeze_batchnorm_statistics(model) | |
| ids = list(train_ids) | |
| rng.shuffle(ids) | |
| total = 0. | |
| for sample_id in ids: | |
| budget() | |
| loaded = load_sample(config.root, training[sample_id], train.expected_hashes) | |
| example = transform_example(convert_segmentation_sample(loaded), horizontal=False if diagnostic else rng.random() < .5, | |
| vertical=False if diagnostic else rng.random() < .5, rotations=0 if diagnostic else rng.randrange(4)) | |
| if example.image.shape != (1024, 1024, 3): | |
| raise ValueError("native 1024 RGB inputs required") | |
| batch = {key: value.to(device) if hasattr(value, "to") else value | |
| for key, value in segmentation_batch([example]).items()} | |
| optimizer.zero_grad(set_to_none=True) | |
| loss, _ = model.criterion(model(batch["img"]), batch) | |
| if not torch.isfinite(loss): | |
| raise FloatingPointError("nonfinite semantic loss") | |
| loss.backward() | |
| optimizer.step() | |
| total += loss.detach().item() | |
| budget() | |
| metrics = evaluate() | |
| if config.freeze_bn_statistics and any( | |
| not torch.equal(value, batchnorm_buffers(model)[name]) | |
| for name, value in initial_bn.items()): | |
| raise RuntimeError("frozen BatchNorm buffers changed") | |
| history.append({"epoch": epoch + 1, | |
| "bn_buffers_unchanged": True if config.freeze_bn_statistics else None, "mean_loss": total / len(ids), "metrics": metrics, | |
| "elapsed_seconds": time.monotonic() - started}) | |
| if metrics["miou"] > best: | |
| best, stale, best_metrics = metrics["miou"], 0, metrics | |
| best_checkpoint = str(output / "best.pt") | |
| temporary = output / "best.pt.partial" | |
| torch.save({"model": deepcopy(model).cpu(), "optimizer": optimizer.state_dict(), | |
| "train_args": vars(model.args).copy(), "epoch": epoch + 1, | |
| "augmentation": {"horizontal_flip_probability": 0 if diagnostic else .5, "vertical_flip_probability": 0 if diagnostic else .5, | |
| "uniform_quarter_turns": [0] if diagnostic else [0, 1, 2, 3], "photometric": False, "scale_jitter": False}, | |
| "diagnostic_train_only": diagnostic, "initial_state_sha256": initial_sha, | |
| "metrics": metrics, "transfer": transfer, | |
| "train_ids": train_ids, "validation_ids": validation_ids, | |
| "config": {k: str(v) if isinstance(v, Path) else v for k, v in vars(config).items()}, | |
| "corpus_content_sha256": json.loads((config.root / "stats.json").read_text())["validation_content_sha256"]}, temporary) | |
| temporary.replace(best_checkpoint) | |
| else: | |
| stale += 1 | |
| progress = {"status": "IN_PROGRESS", "best_checkpoint": best_checkpoint, | |
| "history": history, "best_metrics": best_metrics, | |
| "elapsed_seconds": time.monotonic() - started} | |
| temporary_progress = output / "progress.json.partial" | |
| temporary_progress.write_text(json.dumps(progress, indent=2), encoding="utf-8") | |
| temporary_progress.replace(output / "progress.json") | |
| if stale >= config.patience: | |
| status = "EARLY_STOPPED" | |
| break | |
| except TimeoutError: | |
| status = "TIME_LIMIT" | |
| except Exception as error: | |
| (output / "fit_STOP.json").write_text(json.dumps({"status": "FAILED", "diagnosis": f"{type(error).__name__}: {error}"}), encoding="utf-8") | |
| raise | |
| result = {"status": status, "best_checkpoint": best_checkpoint, "history": history, | |
| "best_metrics": best_metrics, "elapsed_seconds": time.monotonic() - started} | |
| (output / "fit.json").write_text(json.dumps(result, indent=2), encoding="utf-8") | |
| return result | |
| def make_segmentation_trainer(config: SegmenterTrainConfig) -> Trainer: | |
| return Trainer(config) | |
Xet Storage Details
- Size:
- 13.7 kB
- Xet hash:
- aadab3ef746258a8fd270065b3be77fc631ef9266e8f5cf5124b2ec6e192d081
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.