Story 8
This commit is contained in:
Ryan
@@ -0,0 +1 @@
|
||||
"""Active learning pipeline: label ingestion, training, validation, and deployment."""
|
||||
@@ -0,0 +1,217 @@
|
||||
"""
|
||||
Label ingestion for active learning.
|
||||
|
||||
Queries the review_queue for annotated labels, extracts face crops,
|
||||
and builds a versioned training dataset with stratified train/val split.
|
||||
"""
|
||||
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import random
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Train split fraction (val = 1 - TRAIN_FRAC)
|
||||
_TRAIN_FRAC = 0.80
|
||||
|
||||
|
||||
class LabelIngestor:
|
||||
"""Ingest annotated review data and build a versioned dataset."""
|
||||
|
||||
def __init__(self, db_connector, config):
|
||||
self._db = db_connector
|
||||
self._cfg = config
|
||||
self._training_path = Path(config.get("storage.training_path", "/data/training"))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def count_annotated(self) -> int:
|
||||
"""Return number of annotated rows ready for training."""
|
||||
rows = self._db.fetchall(
|
||||
"SELECT COUNT(*) AS cnt FROM review_queue "
|
||||
"WHERE annotated = TRUE AND ground_truth IS NOT NULL"
|
||||
)
|
||||
return int((rows or [{"cnt": 0}])[0]["cnt"])
|
||||
|
||||
def ingest(self, version: str, seed: int = 42) -> Optional[str]:
|
||||
"""
|
||||
Build a versioned dataset directory from the review queue.
|
||||
|
||||
Returns the dataset path on success, or None if no data is available.
|
||||
"""
|
||||
rows = self._fetch_annotated_rows()
|
||||
if not rows:
|
||||
logger.warning("No annotated samples found; skipping dataset build")
|
||||
return None
|
||||
|
||||
dataset_path = self._training_path / version
|
||||
crops_pos = dataset_path / "crops" / "class_1"
|
||||
crops_neg = dataset_path / "crops" / "class_0"
|
||||
for d in (crops_pos, crops_neg):
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
records = self._extract_crops(rows, crops_pos, crops_neg)
|
||||
if not records:
|
||||
logger.error("No crops could be extracted from annotated rows")
|
||||
return None
|
||||
|
||||
train_records, val_records = _stratified_split(records, _TRAIN_FRAC, seed)
|
||||
|
||||
self._write_labels_csv(dataset_path / "labels.csv", records)
|
||||
self._write_metadata(dataset_path / "metadata.json", version, records, train_records, val_records)
|
||||
|
||||
logger.info(
|
||||
"Dataset v%s built: %d total (%d train / %d val) at %s",
|
||||
version, len(records), len(train_records), len(val_records), dataset_path,
|
||||
)
|
||||
return str(dataset_path)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _fetch_annotated_rows(self) -> List[Dict]:
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT rq.video_id, rq.ground_truth,
|
||||
pl.confidence_scores,
|
||||
v.file_path
|
||||
FROM review_queue rq
|
||||
JOIN videos v ON v.id = rq.video_id
|
||||
LEFT JOIN (
|
||||
SELECT video_id, confidence_scores,
|
||||
ROW_NUMBER() OVER (PARTITION BY video_id ORDER BY processed_at DESC) AS rn
|
||||
FROM processing_logs
|
||||
) pl ON pl.video_id = rq.video_id AND pl.rn = 1
|
||||
WHERE rq.annotated = TRUE AND rq.ground_truth IS NOT NULL"""
|
||||
)
|
||||
return rows or []
|
||||
|
||||
def _extract_crops(
|
||||
self,
|
||||
rows: List[Dict],
|
||||
crops_pos: Path,
|
||||
crops_neg: Path,
|
||||
) -> List[Dict]:
|
||||
"""Copy or re-extract face crops into the dataset directory."""
|
||||
records: List[Dict] = []
|
||||
scratch_path = Path(self._cfg.get("storage.scratch_path", "/scratch"))
|
||||
|
||||
for row in rows:
|
||||
video_id = row["video_id"]
|
||||
label = int(bool(row["ground_truth"]))
|
||||
dest_dir = crops_pos if label == 1 else crops_neg
|
||||
|
||||
try:
|
||||
contributing = _parse_json_field(row.get("confidence_scores"))
|
||||
crop_paths = self._resolve_crop_paths(video_id, contributing, scratch_path)
|
||||
|
||||
for crop_path in crop_paths:
|
||||
src = Path(crop_path)
|
||||
if not src.exists():
|
||||
logger.debug("Crop not found, skipping: %s", src)
|
||||
continue
|
||||
dest = dest_dir / f"video{video_id}_{src.name}"
|
||||
shutil.copy2(src, dest)
|
||||
records.append({"crop_path": str(dest), "label": label, "video_id": video_id})
|
||||
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to extract crops for video %s: %s", video_id, exc)
|
||||
|
||||
return records
|
||||
|
||||
def _resolve_crop_paths(
|
||||
self,
|
||||
video_id: int,
|
||||
contributing_frames: List,
|
||||
scratch_path: Path,
|
||||
) -> List[str]:
|
||||
"""Return paths of existing crop files for this video."""
|
||||
paths: List[str] = []
|
||||
|
||||
# contributing_frames may be a list of frame scores or frame dicts
|
||||
if contributing_frames:
|
||||
for item in contributing_frames:
|
||||
if isinstance(item, dict) and "crop_path" in item:
|
||||
paths.append(item["crop_path"])
|
||||
elif isinstance(item, str):
|
||||
paths.append(item)
|
||||
|
||||
# Fall back: glob any crop files already written to scratch
|
||||
if not paths:
|
||||
pattern = f"*video{video_id}*.jpg"
|
||||
paths = [str(p) for p in scratch_path.rglob(pattern)]
|
||||
|
||||
return paths
|
||||
|
||||
@staticmethod
|
||||
def _write_labels_csv(csv_path: Path, records: List[Dict]):
|
||||
with open(csv_path, "w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=["crop_path", "label", "video_id"])
|
||||
writer.writeheader()
|
||||
writer.writerows(records)
|
||||
|
||||
@staticmethod
|
||||
def _write_metadata(
|
||||
meta_path: Path,
|
||||
version: str,
|
||||
all_records: List[Dict],
|
||||
train_records: List[Dict],
|
||||
val_records: List[Dict],
|
||||
):
|
||||
pos = sum(1 for r in all_records if r["label"] == 1)
|
||||
neg = len(all_records) - pos
|
||||
meta = {
|
||||
"version": version,
|
||||
"total_samples": len(all_records),
|
||||
"train_samples": len(train_records),
|
||||
"val_samples": len(val_records),
|
||||
"class_counts": {"class_0": neg, "class_1": pos},
|
||||
}
|
||||
meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Module-level helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _stratified_split(
|
||||
records: List[Dict],
|
||||
train_frac: float,
|
||||
seed: int,
|
||||
) -> Tuple[List[Dict], List[Dict]]:
|
||||
"""Return (train, val) with stratification by label."""
|
||||
rng = random.Random(seed)
|
||||
|
||||
by_label: Dict[int, List[Dict]] = {}
|
||||
for r in records:
|
||||
by_label.setdefault(r["label"], []).append(r)
|
||||
|
||||
train, val = [], []
|
||||
for label_records in by_label.values():
|
||||
shuffled = list(label_records)
|
||||
rng.shuffle(shuffled)
|
||||
split = max(1, int(len(shuffled) * train_frac))
|
||||
train.extend(shuffled[:split])
|
||||
val.extend(shuffled[split:])
|
||||
|
||||
rng.shuffle(train)
|
||||
rng.shuffle(val)
|
||||
return train, val
|
||||
|
||||
|
||||
def _parse_json_field(value) -> List:
|
||||
if not value:
|
||||
return []
|
||||
if isinstance(value, list):
|
||||
return value
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
return parsed if isinstance(parsed, list) else []
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return []
|
||||
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
Active learning pipeline orchestrator.
|
||||
|
||||
Entry point that coordinates:
|
||||
1. Check annotated sample count against minimum threshold
|
||||
2. Ingest labels and build versioned dataset
|
||||
3. Fine-tune classification head
|
||||
4. Validate candidate model against quality gates
|
||||
5. Deploy if gates pass; skip otherwise
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from .label_ingestor import LabelIngestor
|
||||
from .registry import ModelRegistry
|
||||
from .trainer import Trainer
|
||||
from .validator import Validator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ActiveLearningPipeline:
|
||||
"""Orchestrate the full active learning cycle."""
|
||||
|
||||
def __init__(self, db_connector, config):
|
||||
self._db = db_connector
|
||||
self._cfg = config
|
||||
self._ingestor = LabelIngestor(db_connector, config)
|
||||
self._registry = ModelRegistry(db_connector, config)
|
||||
|
||||
def run(self, version: str) -> bool:
|
||||
"""
|
||||
Execute a full active learning cycle for the given model version.
|
||||
|
||||
Returns True if a new model was deployed, False otherwise.
|
||||
"""
|
||||
al = self._cfg.get_section("active_learning")
|
||||
min_samples: int = int(al.get("min_annotated_samples", 100))
|
||||
|
||||
# 1. Check threshold
|
||||
n_annotated = self._ingestor.count_annotated()
|
||||
logger.info("Annotated samples available: %d (minimum: %d)", n_annotated, min_samples)
|
||||
if n_annotated < min_samples:
|
||||
logger.info("Insufficient annotated samples; skipping active learning cycle")
|
||||
return False
|
||||
|
||||
# 2. Ingest labels
|
||||
seed = int(self._cfg.get("active_learning.seed", 42))
|
||||
dataset_path = self._ingestor.ingest(version, seed=seed)
|
||||
if not dataset_path:
|
||||
logger.error("Label ingestion produced no dataset; aborting")
|
||||
return False
|
||||
|
||||
# 3. Train
|
||||
trainer = Trainer(self._cfg)
|
||||
checkpoint_path = trainer.train(dataset_path, version)
|
||||
if not checkpoint_path:
|
||||
logger.error("Training failed; aborting active learning cycle")
|
||||
return False
|
||||
|
||||
# 4. Validate
|
||||
current_f1 = self._registry.get_active_f1()
|
||||
validator = Validator(self._cfg, current_f1=current_f1)
|
||||
metrics = validator.validate(checkpoint_path, dataset_path)
|
||||
|
||||
self._registry.register_candidate(
|
||||
version, checkpoint_path,
|
||||
f1=metrics.get("f1", 0.0),
|
||||
ece=metrics.get("ece", 1.0),
|
||||
)
|
||||
|
||||
if not metrics.get("gates_passed", False):
|
||||
logger.warning(
|
||||
"Quality gates failed for %s; skipping deployment. Details: %s",
|
||||
version, metrics.get("gate_details"),
|
||||
)
|
||||
return False
|
||||
|
||||
# 5. Deploy
|
||||
deployed = self._registry.deploy(version, checkpoint_path)
|
||||
return deployed
|
||||
@@ -0,0 +1,223 @@
|
||||
"""
|
||||
Model registry and deployment for active learning.
|
||||
|
||||
Handles:
|
||||
- Promoting a candidate model to ACTIVE in the DB
|
||||
- Converting the PyTorch checkpoint to ONNX then TensorRT
|
||||
- Hot-reloading the classifier in worker processes via SIGUSER1
|
||||
- Rollback to the previous ACTIVE model on failure
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Signal used to notify the worker process to reload its model
|
||||
_RELOAD_SIGNAL = signal.SIGUSR1
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""Manage model lifecycle: promote, deploy, rollback."""
|
||||
|
||||
def __init__(self, db_connector, config):
|
||||
self._db = db_connector
|
||||
self._cfg = config
|
||||
self._models_path = Path(config.get("storage.models_path", "/models"))
|
||||
al = config.get_section("active_learning")
|
||||
dep = al.get("deployment", {})
|
||||
self._auto_deploy: bool = bool(dep.get("auto_deploy", True))
|
||||
self._hot_reload: bool = bool(dep.get("hot_reload", True))
|
||||
self._rollback_enabled: bool = bool(dep.get("rollback_enabled", True))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_active_version(self) -> Optional[str]:
|
||||
"""Return the version string of the currently ACTIVE model."""
|
||||
rows = self._db.fetchall(
|
||||
"SELECT version FROM models WHERE status = 'ACTIVE' ORDER BY deployed_at DESC LIMIT 1"
|
||||
)
|
||||
if rows:
|
||||
return rows[0]["version"]
|
||||
return None
|
||||
|
||||
def get_active_f1(self) -> float:
|
||||
"""Return the F1 score of the currently ACTIVE model (0 if unknown)."""
|
||||
rows = self._db.fetchall(
|
||||
"SELECT f1_score FROM models WHERE status = 'ACTIVE' ORDER BY deployed_at DESC LIMIT 1"
|
||||
)
|
||||
if rows and rows[0]["f1_score"] is not None:
|
||||
return float(rows[0]["f1_score"])
|
||||
return 0.0
|
||||
|
||||
def register_candidate(self, version: str, checkpoint_path: str, f1: float, ece: float):
|
||||
"""Insert or update the candidate model record."""
|
||||
self._db.execute(
|
||||
"""INSERT INTO models (version, status, path, f1_score, ece_score, created_at)
|
||||
VALUES (%s, 'CANDIDATE', %s, %s, %s, NOW())
|
||||
ON DUPLICATE KEY UPDATE
|
||||
status = 'CANDIDATE', path = VALUES(path),
|
||||
f1_score = VALUES(f1_score), ece_score = VALUES(ece_score)""",
|
||||
(version, checkpoint_path, f1, ece),
|
||||
)
|
||||
logger.info("Registered candidate model %s (f1=%.4f, ece=%.4f)", version, f1, ece)
|
||||
|
||||
def deploy(self, version: str, checkpoint_path: str) -> bool:
|
||||
"""
|
||||
Convert, promote and hot-reload the candidate model.
|
||||
|
||||
Returns True if deployment succeeded, False otherwise.
|
||||
"""
|
||||
if not self._auto_deploy:
|
||||
logger.info("Auto-deploy disabled; skipping deployment of %s", version)
|
||||
return False
|
||||
|
||||
try:
|
||||
engine_path = self._build_engine(version, checkpoint_path)
|
||||
self._promote(version, engine_path)
|
||||
if self._hot_reload:
|
||||
self._reload_worker()
|
||||
logger.info("Model %s deployed successfully", version)
|
||||
return True
|
||||
except Exception as exc:
|
||||
logger.error("Deployment of %s failed: %s", version, exc, exc_info=True)
|
||||
if self._rollback_enabled:
|
||||
self.rollback(version)
|
||||
return False
|
||||
|
||||
def rollback(self, failed_version: str):
|
||||
"""Restore the most recent ARCHIVED model to ACTIVE."""
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT version FROM models
|
||||
WHERE status = 'ARCHIVED'
|
||||
ORDER BY deployed_at DESC
|
||||
LIMIT 1"""
|
||||
)
|
||||
if not rows:
|
||||
logger.error("No archived model available for rollback")
|
||||
return
|
||||
|
||||
prev_version = rows[0]["version"]
|
||||
self._db.execute(
|
||||
"UPDATE models SET status = 'ARCHIVED' WHERE version = %s",
|
||||
(failed_version,),
|
||||
)
|
||||
self._db.execute(
|
||||
"UPDATE models SET status = 'ACTIVE' WHERE version = %s",
|
||||
(prev_version,),
|
||||
)
|
||||
logger.warning("Rolled back to model %s (failed: %s)", prev_version, failed_version)
|
||||
if self._hot_reload:
|
||||
self._reload_worker()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_engine(self, version: str, checkpoint_path: str) -> str:
|
||||
"""Export checkpoint to ONNX then build a TensorRT FP32 engine."""
|
||||
candidate_dir = self._models_path / "candidate"
|
||||
onnx_path = candidate_dir / f"{version}.onnx"
|
||||
engine_path = candidate_dir / f"{version}.trt"
|
||||
|
||||
self._export_onnx(checkpoint_path, onnx_path)
|
||||
self._build_trt(onnx_path, engine_path)
|
||||
return str(engine_path)
|
||||
|
||||
@staticmethod
|
||||
def _export_onnx(checkpoint_path: str, onnx_path: Path):
|
||||
"""Convert PyTorch checkpoint to ONNX."""
|
||||
try:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.models as models
|
||||
|
||||
device = torch.device("cpu")
|
||||
model = models.mobilenet_v3_small(weights=None)
|
||||
model.classifier = nn.Sequential(
|
||||
nn.Linear(576, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 2),
|
||||
)
|
||||
state = torch.load(checkpoint_path, map_location=device)
|
||||
model.load_state_dict(state)
|
||||
model.eval()
|
||||
|
||||
dummy = torch.zeros(1, 3, 224, 224)
|
||||
torch.onnx.export(
|
||||
model,
|
||||
dummy,
|
||||
str(onnx_path),
|
||||
input_names=["input"],
|
||||
output_names=["output"],
|
||||
dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
|
||||
opset_version=11,
|
||||
)
|
||||
logger.info("Exported ONNX model to %s", onnx_path)
|
||||
except ImportError as exc:
|
||||
logger.error("PyTorch/torchvision not available for ONNX export: %s", exc)
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _build_trt(onnx_path: Path, engine_path: Path):
|
||||
"""Build TensorRT FP32 engine via trtexec CLI."""
|
||||
cmd = [
|
||||
"trtexec",
|
||||
f"--onnx={onnx_path}",
|
||||
f"--saveEngine={engine_path}",
|
||||
"--fp32",
|
||||
"--maxBatch=32",
|
||||
"--workspace=2048",
|
||||
]
|
||||
logger.info("Building TensorRT engine: %s", " ".join(cmd))
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"trtexec failed (rc={result.returncode}): {result.stderr[:500]}"
|
||||
)
|
||||
logger.info("TensorRT engine saved to %s", engine_path)
|
||||
|
||||
def _promote(self, version: str, engine_path: str):
|
||||
"""Atomically promote candidate to ACTIVE, archive current ACTIVE."""
|
||||
self._db.execute(
|
||||
"""UPDATE models
|
||||
SET status = 'ARCHIVED'
|
||||
WHERE status = 'ACTIVE'""",
|
||||
)
|
||||
self._db.execute(
|
||||
"""UPDATE models
|
||||
SET status = 'ACTIVE', path = %s, deployed_at = NOW()
|
||||
WHERE version = %s""",
|
||||
(engine_path, version),
|
||||
)
|
||||
logger.info("Promoted model %s to ACTIVE (engine=%s)", version, engine_path)
|
||||
|
||||
@staticmethod
|
||||
def _reload_worker():
|
||||
"""Send SIGUSR1 to the worker process to trigger a model hot-reload."""
|
||||
pid_file = Path("/tmp/videodetect_worker.pid")
|
||||
if pid_file.exists():
|
||||
try:
|
||||
pid = int(pid_file.read_text().strip())
|
||||
os.kill(pid, _RELOAD_SIGNAL)
|
||||
logger.info("Sent SIGUSR1 to worker PID %d for hot-reload", pid)
|
||||
except (ValueError, ProcessLookupError, PermissionError) as exc:
|
||||
logger.warning("Could not signal worker process: %s", exc)
|
||||
else:
|
||||
# Fallback: restart via docker-compose
|
||||
logger.warning("No PID file found; attempting docker-compose restart")
|
||||
try:
|
||||
subprocess.run(
|
||||
["docker-compose", "restart", "worker"],
|
||||
capture_output=True,
|
||||
timeout=60,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("docker-compose restart failed: %s", exc)
|
||||
@@ -0,0 +1,344 @@
|
||||
"""
|
||||
Training pipeline for active learning head-only fine-tuning.
|
||||
|
||||
Loads a pre-trained MobileNetV3 backbone (frozen), replaces the classification
|
||||
head, and fine-tunes on versioned crop datasets using AdamW + early stopping.
|
||||
All training is FP32 for Tesla P40 compatibility (CC 5.2, no Tensor Cores).
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Trainer:
|
||||
"""Head-only fine-tuning pipeline for MobileNetV3 classifier."""
|
||||
|
||||
def __init__(self, config):
|
||||
self._cfg = config
|
||||
self._models_path = Path(config.get("storage.models_path", "/models"))
|
||||
self._candidate_path = self._models_path / "candidate"
|
||||
self._candidate_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
al = config.get_section("active_learning")
|
||||
tr = al.get("training", {})
|
||||
self._epochs: int = int(tr.get("epochs", 20))
|
||||
self._batch_size: int = int(tr.get("batch_size", 32))
|
||||
self._lr: float = float(tr.get("learning_rate", 1e-3))
|
||||
self._weight_decay: float = float(tr.get("weight_decay", 1e-2))
|
||||
self._patience: int = int(tr.get("early_stopping_patience", 5))
|
||||
self._lr_factor: float = float(tr.get("lr_factor", 0.5))
|
||||
self._lr_patience: int = int(tr.get("lr_patience", 3))
|
||||
|
||||
aug = al.get("augmentation", {})
|
||||
self._aug_flip: bool = bool(aug.get("horizontal_flip", True))
|
||||
self._aug_jitter: bool = bool(aug.get("color_jitter", True))
|
||||
self._aug_affine: bool = bool(aug.get("affine", True))
|
||||
self._aug_degrees: int = int(aug.get("affine_degrees", 10))
|
||||
self._aug_scale: float = float(aug.get("affine_scale", 0.1))
|
||||
|
||||
seed = int(config.get("active_learning.seed", 42))
|
||||
self._seed = seed
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def train(self, dataset_path: str, version: str) -> Optional[str]:
|
||||
"""
|
||||
Fine-tune the classification head on the versioned dataset.
|
||||
|
||||
Returns the path to the best checkpoint or None on failure.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
self._seed_everything(self._seed)
|
||||
self._assert_cuda_compatibility(torch)
|
||||
|
||||
train_loader, val_loader = self._build_dataloaders(dataset_path, torch)
|
||||
model = self._build_model(nn)
|
||||
best_path = self._run_training(model, train_loader, val_loader, version, torch, nn)
|
||||
return best_path
|
||||
except ImportError as exc:
|
||||
logger.error("PyTorch not available: %s", exc)
|
||||
return None
|
||||
except Exception as exc:
|
||||
logger.error("Training failed: %s", exc, exc_info=True)
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Model construction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_model(self, nn):
|
||||
"""Load MobileNetV3-Small backbone, freeze it, attach trainable head."""
|
||||
import torchvision.models as models
|
||||
|
||||
model = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1)
|
||||
|
||||
# Freeze backbone
|
||||
for param in model.features.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
# Replace classification head
|
||||
model.classifier = nn.Sequential(
|
||||
nn.Linear(576, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 2),
|
||||
)
|
||||
return model
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Data loading
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_dataloaders(self, dataset_path: str, torch) -> Tuple:
|
||||
from torch.utils.data import DataLoader
|
||||
import torchvision.transforms as T
|
||||
|
||||
train_transform = self._build_train_transform(T)
|
||||
val_transform = self._build_val_transform(T)
|
||||
|
||||
train_dataset = _CropDataset(
|
||||
Path(dataset_path) / "crops", split="train",
|
||||
transform=train_transform, seed=self._seed,
|
||||
)
|
||||
val_dataset = _CropDataset(
|
||||
Path(dataset_path) / "crops", split="val",
|
||||
transform=val_transform, seed=self._seed,
|
||||
)
|
||||
|
||||
g = torch.Generator()
|
||||
g.manual_seed(self._seed)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=self._batch_size,
|
||||
shuffle=True,
|
||||
num_workers=2,
|
||||
pin_memory=True,
|
||||
generator=g,
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
val_dataset,
|
||||
batch_size=self._batch_size,
|
||||
shuffle=False,
|
||||
num_workers=2,
|
||||
pin_memory=True,
|
||||
)
|
||||
logger.info(
|
||||
"Dataset: %d train / %d val samples",
|
||||
len(train_dataset), len(val_dataset),
|
||||
)
|
||||
return train_loader, val_loader
|
||||
|
||||
def _build_train_transform(self, T):
|
||||
ops = [T.Resize((224, 224))]
|
||||
if self._aug_flip:
|
||||
ops.append(T.RandomHorizontalFlip())
|
||||
if self._aug_jitter:
|
||||
ops.append(T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1))
|
||||
if self._aug_affine:
|
||||
ops.append(T.RandomAffine(
|
||||
degrees=self._aug_degrees,
|
||||
scale=(1.0 - self._aug_scale, 1.0 + self._aug_scale),
|
||||
))
|
||||
ops += [T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
|
||||
return T.Compose(ops)
|
||||
|
||||
def _build_val_transform(self, T):
|
||||
return T.Compose([
|
||||
T.Resize((224, 224)),
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Training loop
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_training(self, model, train_loader, val_loader, version: str, torch, nn) -> Optional[str]:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
model = model.to(device)
|
||||
logger.info("Training on device: %s", device)
|
||||
|
||||
class_weights = self._compute_class_weights(train_loader, torch, device)
|
||||
criterion = nn.CrossEntropyLoss(weight=class_weights)
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
filter(lambda p: p.requires_grad, model.parameters()),
|
||||
lr=self._lr,
|
||||
weight_decay=self._weight_decay,
|
||||
betas=(0.9, 0.999),
|
||||
)
|
||||
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
||||
optimizer, mode="max", factor=self._lr_factor, patience=self._lr_patience,
|
||||
)
|
||||
|
||||
best_f1 = -1.0
|
||||
best_path: Optional[str] = None
|
||||
no_improve = 0
|
||||
|
||||
for epoch in range(1, self._epochs + 1):
|
||||
train_loss = self._train_epoch(model, train_loader, criterion, optimizer, device, torch)
|
||||
val_f1 = self._val_epoch(model, val_loader, device, torch)
|
||||
scheduler.step(val_f1)
|
||||
|
||||
ckpt_path = self._candidate_path / f"{version}_epoch_{epoch:02d}.pt"
|
||||
torch.save(model.state_dict(), ckpt_path)
|
||||
logger.info(
|
||||
"Epoch %d/%d train_loss=%.4f val_f1=%.4f",
|
||||
epoch, self._epochs, train_loss, val_f1,
|
||||
)
|
||||
|
||||
if val_f1 > best_f1:
|
||||
best_f1 = val_f1
|
||||
best_path = str(self._candidate_path / f"{version}_best.pt")
|
||||
torch.save(model.state_dict(), best_path)
|
||||
no_improve = 0
|
||||
else:
|
||||
no_improve += 1
|
||||
if no_improve >= self._patience:
|
||||
logger.info("Early stopping at epoch %d (patience=%d)", epoch, self._patience)
|
||||
break
|
||||
|
||||
logger.info("Training complete. Best val_f1=%.4f checkpoint=%s", best_f1, best_path)
|
||||
return best_path
|
||||
|
||||
@staticmethod
|
||||
def _train_epoch(model, loader, criterion, optimizer, device, torch) -> float:
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
for images, labels in loader:
|
||||
images, labels = images.to(device), labels.to(device)
|
||||
optimizer.zero_grad()
|
||||
outputs = model(images)
|
||||
loss = criterion(outputs, labels)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
total_loss += loss.item() * len(images)
|
||||
return total_loss / max(len(loader.dataset), 1)
|
||||
|
||||
@staticmethod
|
||||
def _val_epoch(model, loader, device, torch) -> float:
|
||||
from sklearn.metrics import f1_score
|
||||
|
||||
model.eval()
|
||||
all_preds, all_labels = [], []
|
||||
with torch.no_grad():
|
||||
for images, labels in loader:
|
||||
images = images.to(device)
|
||||
outputs = model(images)
|
||||
preds = outputs.argmax(dim=1).cpu().numpy()
|
||||
all_preds.extend(preds)
|
||||
all_labels.extend(labels.numpy())
|
||||
|
||||
return float(f1_score(all_labels, all_preds, average="macro", zero_division=0))
|
||||
|
||||
@staticmethod
|
||||
def _compute_class_weights(train_loader, torch, device):
|
||||
"""Compute inverse-frequency class weights to handle imbalance."""
|
||||
counts = np.zeros(2, dtype=np.int64)
|
||||
for _, labels in train_loader:
|
||||
for lbl in labels.numpy():
|
||||
counts[int(lbl)] += 1
|
||||
total = counts.sum()
|
||||
if total == 0 or counts.min() == 0:
|
||||
return None
|
||||
weights = total / (2.0 * counts.astype(np.float32))
|
||||
return torch.tensor(weights, dtype=torch.float32, device=device)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Misc
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _assert_cuda_compatibility(torch):
|
||||
if torch.cuda.is_available():
|
||||
cuda_ver = torch.version.cuda or ""
|
||||
torch_ver = torch.__version__
|
||||
logger.info("CUDA version: %s PyTorch: %s", cuda_ver, torch_ver)
|
||||
major = int(cuda_ver.split(".")[0]) if cuda_ver else 0
|
||||
if major > 11:
|
||||
logger.warning(
|
||||
"CUDA %s detected; system requires CUDA ≤ 11.8 for Tesla P40", cuda_ver
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _seed_everything(seed: int):
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
try:
|
||||
import torch
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Dataset helper
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
class _CropDataset:
|
||||
"""Reads face crop images from class_0/ and class_1/ subdirectories."""
|
||||
|
||||
_EXTENSIONS = {".jpg", ".jpeg", ".png"}
|
||||
|
||||
def __init__(self, crops_root: Path, split: str, transform, seed: int):
|
||||
all_samples = self._discover(crops_root)
|
||||
train_samples, val_samples = _split_samples(all_samples, 0.80, seed)
|
||||
self._samples = train_samples if split == "train" else val_samples
|
||||
self._transform = transform
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._samples)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
from PIL import Image
|
||||
path, label = self._samples[idx]
|
||||
image = Image.open(path).convert("RGB")
|
||||
if self._transform:
|
||||
image = self._transform(image)
|
||||
return image, label
|
||||
|
||||
def _discover(self, crops_root: Path):
|
||||
samples = []
|
||||
for label, subdir in ((0, "class_0"), (1, "class_1")):
|
||||
d = crops_root / subdir
|
||||
if d.is_dir():
|
||||
for p in sorted(d.iterdir()):
|
||||
if p.suffix.lower() in self._EXTENSIONS:
|
||||
samples.append((p, label))
|
||||
return samples
|
||||
|
||||
|
||||
def _split_samples(samples, train_frac: float, seed: int):
|
||||
"""Stratified split mirroring LabelIngestor logic."""
|
||||
rng = random.Random(seed)
|
||||
by_label: Dict[int, list] = {}
|
||||
for path, label in samples:
|
||||
by_label.setdefault(label, []).append((path, label))
|
||||
|
||||
train, val = [], []
|
||||
for group in by_label.values():
|
||||
shuffled = list(group)
|
||||
rng.shuffle(shuffled)
|
||||
split = max(1, int(len(shuffled) * train_frac))
|
||||
train.extend(shuffled[:split])
|
||||
val.extend(shuffled[split:])
|
||||
|
||||
return train, val
|
||||
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
Candidate model validation for active learning.
|
||||
|
||||
Evaluates a candidate checkpoint on the held-out validation set,
|
||||
computes F1, ECE, accuracy, precision, and recall, then checks
|
||||
quality gates before promotion.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Validator:
|
||||
"""Evaluate a candidate model and enforce quality gates."""
|
||||
|
||||
def __init__(self, config, current_f1: float = 0.0):
|
||||
al = config.get_section("active_learning")
|
||||
val_cfg = al.get("validation", {})
|
||||
self._min_f1_delta: float = float(val_cfg.get("min_f1_improvement", 0.02))
|
||||
self._max_ece: float = float(val_cfg.get("max_ece", 0.08))
|
||||
self._current_f1 = current_f1
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def validate(self, checkpoint_path: str, dataset_path: str) -> Dict:
|
||||
"""
|
||||
Run validation and return a results dict including gate pass/fail.
|
||||
|
||||
Keys: f1, ece, accuracy, precision, recall, gates_passed, gate_details
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
probs, labels = self._run_inference(checkpoint_path, dataset_path, torch)
|
||||
return self._compute_metrics(probs, labels)
|
||||
except ImportError as exc:
|
||||
logger.error("PyTorch not available: %s", exc)
|
||||
return self._empty_result()
|
||||
except Exception as exc:
|
||||
logger.error("Validation failed: %s", exc, exc_info=True)
|
||||
return self._empty_result()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Inference
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_inference(self, checkpoint_path: str, dataset_path: str, torch) -> Tuple[np.ndarray, np.ndarray]:
|
||||
import torch.nn as nn
|
||||
import torchvision.models as models
|
||||
import torchvision.transforms as T
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model = self._load_checkpoint(checkpoint_path, nn, models, torch, device)
|
||||
model.eval()
|
||||
|
||||
transform = T.Compose([
|
||||
T.Resize((224, 224)),
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
from .trainer import _CropDataset, _split_samples
|
||||
crops_root = Path(dataset_path) / "crops"
|
||||
all_samples = _CropDataset._discover(None, crops_root) # type: ignore[arg-type]
|
||||
_, val_samples = _split_samples(all_samples, 0.80, seed=42)
|
||||
|
||||
if not val_samples:
|
||||
raise ValueError("Validation set is empty")
|
||||
|
||||
val_dataset = _ValDataset(val_samples, transform)
|
||||
loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=2, pin_memory=True)
|
||||
|
||||
all_probs, all_labels = [], []
|
||||
with torch.no_grad():
|
||||
for images, batch_labels in loader:
|
||||
images = images.to(device)
|
||||
outputs = model(images)
|
||||
probs_batch = torch.softmax(outputs, dim=1)[:, 1].cpu().numpy()
|
||||
all_probs.extend(probs_batch)
|
||||
all_labels.extend(batch_labels.numpy())
|
||||
|
||||
return np.array(all_probs, dtype=np.float32), np.array(all_labels, dtype=np.int32)
|
||||
|
||||
@staticmethod
|
||||
def _load_checkpoint(checkpoint_path: str, nn, models, torch, device):
|
||||
model = models.mobilenet_v3_small(weights=None)
|
||||
model.classifier = nn.Sequential(
|
||||
nn.Linear(576, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 2),
|
||||
)
|
||||
state = torch.load(checkpoint_path, map_location=device)
|
||||
model.load_state_dict(state)
|
||||
return model.to(device)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Metrics
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _compute_metrics(self, probs: np.ndarray, labels: np.ndarray) -> Dict:
|
||||
from sklearn.metrics import (
|
||||
accuracy_score, f1_score, precision_score, recall_score,
|
||||
)
|
||||
|
||||
preds = (probs >= 0.5).astype(np.int32)
|
||||
f1 = float(f1_score(labels, preds, average="macro", zero_division=0))
|
||||
accuracy = float(accuracy_score(labels, preds))
|
||||
precision = float(precision_score(labels, preds, average="macro", zero_division=0))
|
||||
recall = float(recall_score(labels, preds, average="macro", zero_division=0))
|
||||
ece = compute_ece(probs, labels)
|
||||
|
||||
delta_f1 = f1 - self._current_f1
|
||||
gate_f1 = delta_f1 > self._min_f1_delta
|
||||
gate_ece = ece < self._max_ece
|
||||
gates_passed = gate_f1 and gate_ece
|
||||
|
||||
gate_details = {
|
||||
"delta_f1": round(delta_f1, 4),
|
||||
"min_f1_improvement": self._min_f1_delta,
|
||||
"gate_f1_passed": gate_f1,
|
||||
"ece": round(ece, 4),
|
||||
"max_ece": self._max_ece,
|
||||
"gate_ece_passed": gate_ece,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Validation results: f1=%.4f delta_f1=%.4f ece=%.4f gates_passed=%s",
|
||||
f1, delta_f1, ece, gates_passed,
|
||||
)
|
||||
|
||||
return {
|
||||
"f1": round(f1, 4),
|
||||
"ece": round(ece, 4),
|
||||
"accuracy": round(accuracy, 4),
|
||||
"precision": round(precision, 4),
|
||||
"recall": round(recall, 4),
|
||||
"gates_passed": gates_passed,
|
||||
"gate_details": gate_details,
|
||||
}
|
||||
|
||||
def _empty_result(self) -> Dict:
|
||||
return {
|
||||
"f1": 0.0, "ece": 1.0, "accuracy": 0.0,
|
||||
"precision": 0.0, "recall": 0.0,
|
||||
"gates_passed": False,
|
||||
"gate_details": {},
|
||||
}
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# ECE calculation (module-level for reuse in tests)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def compute_ece(predictions: np.ndarray, labels: np.ndarray, n_bins: int = 15) -> float:
|
||||
"""
|
||||
Expected Calibration Error.
|
||||
|
||||
Bins predictions by confidence and measures average |confidence - accuracy|
|
||||
weighted by bin population.
|
||||
"""
|
||||
bin_boundaries = np.linspace(0, 1, n_bins + 1)
|
||||
ece = 0.0
|
||||
n = len(predictions)
|
||||
if n == 0:
|
||||
return 0.0
|
||||
|
||||
for i in range(n_bins):
|
||||
mask = (predictions >= bin_boundaries[i]) & (predictions < bin_boundaries[i + 1])
|
||||
if mask.sum() > 0:
|
||||
bin_confidence = float(predictions[mask].mean())
|
||||
bin_accuracy = float(labels[mask].mean())
|
||||
ece += (mask.sum() / n) * abs(bin_confidence - bin_accuracy)
|
||||
|
||||
return float(ece)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Minimal dataset wrapper for validation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
class _ValDataset:
|
||||
def __init__(self, samples, transform):
|
||||
self._samples = samples
|
||||
self._transform = transform
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._samples)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
from PIL import Image
|
||||
path, label = self._samples[idx]
|
||||
image = Image.open(path).convert("RGB")
|
||||
if self._transform:
|
||||
image = self._transform(image)
|
||||
return image, label
|
||||
Reference in New Issue
Block a user