"""Post-training optimization: pruning, distillation, NAS, quantization."""
from __future__ import annotations
from typing import Any
import numpy as np
from arcade._logging import get_logger
from arcade.nn.compression import CompressionPoint
log = get_logger(__name__)
_NN_CLASSIFIERS = frozenset(
{"fnn", "hybrid", "herqules", "herqules_leakage", "leakage", "multilevel", "transformer"},
)
[docs]
def run_optimization_stage(
cfg: Any,
data_out: dict[str, Any],
classify_out: dict[str, Any],
) -> dict[str, Any]:
"""Apply optional compression / NAS / quantization to trained models."""
opt = getattr(cfg, "optimization", None)
tuning_results = _collect_tuning_results(classify_out)
if opt is None or not getattr(opt, "enabled", False):
return {"optimized": {}, "tuning_results": tuning_results, "pareto": {}}
splits = data_out.get("splits")
if splits is None:
log.warning("No splits available — skipping optimization stage")
return {"optimized": {}, "tuning_results": tuning_results, "pareto": {}}
clf_results = classify_out.get("classifiers", {})
targets = set(opt.target_classifiers or []) or None
optimized: dict[str, Any] = {}
pareto: dict[str, list[dict[str, Any]]] = {}
for clf_name, bundle in clf_results.items():
if targets is not None and clf_name not in targets:
continue
if clf_name not in _NN_CLASSIFIERS:
continue
model = bundle.get("model")
if model is None:
continue
feats = bundle.get("benchmark_features")
if feats is None:
continue
orig_acc = float(getattr(bundle.get("metrics"), "accuracy", 0.0) or 0.0)
budget = float(opt.error_budget)
candidates: list[dict[str, Any]] = [
{
"model": model,
"accuracy": orig_acc,
"method": "original",
"n_parameters": _count_params(model),
},
]
if opt.pruning.enabled:
candidates.extend(
_prune_candidates(model, feats, splits, orig_acc, budget, opt.pruning),
)
if opt.distillation.enabled:
candidates.extend(
_distill_candidates(model, feats, splits, orig_acc, budget, opt.distillation),
)
if opt.quantization.enabled:
candidates.extend(
_quantize_candidates(model, feats, splits, orig_acc, budget, opt.quantization),
)
if opt.nas.enabled and clf_name in {"fnn", "hybrid", "transformer"}:
candidates.extend(
_nas_candidates(clf_name, feats, splits, orig_acc, budget, opt.nas, model),
)
frontier = _pareto_frontier(candidates, orig_acc, budget)
pareto[clf_name] = [_point_to_dict(p) for p in frontier]
best = _select_best_within_budget(candidates, orig_acc, budget)
optimized[clf_name] = {
**best,
"original_accuracy": orig_acc,
"pareto": pareto[clf_name],
}
log.info(
"Optimization %s: method=%s accuracy=%.4f (orig=%.4f)",
clf_name,
best.get("method"),
best.get("accuracy", 0.0),
orig_acc,
)
return {"optimized": optimized, "tuning_results": tuning_results, "pareto": pareto}
def _collect_tuning_results(classify_out: dict[str, Any]) -> dict[str, Any]:
out: dict[str, Any] = {}
for name, bundle in classify_out.get("classifiers", {}).items():
params = bundle.get("best_params")
if params:
out[name] = params
return out
def _count_params(model: Any) -> int:
if hasattr(model, "parameters"):
return int(sum(p.numel() for p in model.parameters()))
inner = getattr(model, "_model", None)
if inner is not None and hasattr(inner, "parameters"):
return int(sum(p.numel() for p in inner.parameters()))
return 0
def _pareto_frontier(
candidates: list[dict[str, Any]],
orig_acc: float,
budget: float,
) -> list[CompressionPoint]:
"""Accuracy vs parameter-count Pareto frontier within error budget."""
min_acc = orig_acc - budget
points: list[CompressionPoint] = []
for c in candidates:
acc = float(c.get("accuracy", 0.0))
if acc < min_acc:
continue
points.append(
CompressionPoint(
accuracy=acc,
n_parameters=int(c.get("n_parameters", 0)),
inference_ms=float(c.get("inference_ms", 0.0)),
),
)
if not points:
return []
points.sort(key=lambda p: p.n_parameters)
frontier: list[CompressionPoint] = []
best_acc = -1.0
for pt in points:
if pt.accuracy > best_acc:
frontier.append(pt)
best_acc = pt.accuracy
return frontier
def _select_best_within_budget(
candidates: list[dict[str, Any]],
orig_acc: float,
budget: float,
) -> dict[str, Any]:
"""Pick smallest model that stays within the accuracy budget."""
min_acc = orig_acc - budget
valid = [c for c in candidates if float(c.get("accuracy", 0.0)) >= min_acc]
if not valid:
return candidates[0]
return min(valid, key=lambda c: (int(c.get("n_parameters", 10**9)), -float(c["accuracy"])))
def _point_to_dict(pt: CompressionPoint) -> dict[str, Any]:
return {
"accuracy": pt.accuracy,
"n_parameters": pt.n_parameters,
"inference_ms": pt.inference_ms,
}
def _prune_candidates(
model: Any,
feats: np.ndarray,
splits: Any,
orig_acc: float,
budget: float,
pcfg: Any,
) -> list[dict[str, Any]]:
from arcade.nn.compression import iterative_pruning
out: list[dict[str, Any]] = []
try:
points = iterative_pruning(
model,
feats,
splits.train_labels,
val_features=feats[: min(64, len(feats))],
val_labels=splits.val_labels[: min(64, len(splits.val_labels))]
if splits.val_labels is not None
else None,
prune_fraction_per_step=float(pcfg.prune_fraction_per_step),
n_steps=int(pcfg.n_steps),
finetune_epochs=int(pcfg.finetune_epochs),
accuracy_tolerance=budget,
)
for pt in points:
if pt.accuracy >= orig_acc - budget:
out.append(
{
"model": model,
"accuracy": pt.accuracy,
"method": "pruning",
"n_parameters": pt.n_parameters,
"inference_ms": pt.inference_ms,
},
)
except Exception as exc:
log.warning("Pruning failed for model: %s", exc)
return out
def _distill_candidates(
teacher: Any,
feats: np.ndarray,
splits: Any,
orig_acc: float,
budget: float,
dcfg: Any,
) -> list[dict[str, Any]]:
from arcade.nn.compression import distillation_pareto
out: list[dict[str, Any]] = []
def factory(**kwargs: Any) -> Any:
return teacher
try:
points = distillation_pareto(
teacher,
factory,
feats,
splits.train_labels,
temperature=float(dcfg.temperature),
student_sizes=list(dcfg.student_sizes),
accuracy_tolerance=budget,
)
for pt in points:
if pt.accuracy >= orig_acc - budget:
out.append(
{
"model": teacher,
"accuracy": pt.accuracy,
"method": "distillation",
"n_parameters": pt.n_parameters,
"inference_ms": pt.inference_ms,
},
)
except Exception as exc:
log.warning("Distillation failed: %s", exc)
return out
def _quantize_candidates(
model: Any,
feats: np.ndarray,
splits: Any,
orig_acc: float,
budget: float,
qcfg: Any,
) -> list[dict[str, Any]]:
from arcade.nn.quantization import quantize_model
out: list[dict[str, Any]] = []
try:
result = quantize_model(
model,
feats,
splits.train_labels,
precision=str(qcfg.precision),
finetune_epochs=int(qcfg.finetune_epochs),
)
if result.accuracy >= orig_acc - budget:
out.append(
{
"model": result.model if result.model is not None else model,
"accuracy": result.accuracy,
"method": "quantization",
"n_parameters": result.n_parameters,
"precision": result.precision,
"inference_ms": 0.0,
"model_size_bytes": result.model_size_bytes,
"inference_speedup": result.inference_speedup,
},
)
except Exception as exc:
log.warning("Quantization failed: %s", exc)
return out
def _nas_candidates(
clf_name: str,
feats: np.ndarray,
splits: Any,
orig_acc: float,
budget: float,
ncfg: Any,
model: Any,
) -> list[dict[str, Any]]:
from arcade.nn.nas import search_architecture
out: list[dict[str, Any]] = []
try:
result = search_architecture(
feats,
splits.train_labels,
feats,
splits.val_labels,
n_trials=int(ncfg.n_trials),
max_hidden_layers=int(ncfg.max_hidden_layers),
min_hidden_dim=int(ncfg.min_hidden_dim),
max_hidden_dim=int(ncfg.max_hidden_dim),
hardware_aware=bool(ncfg.hardware_aware),
)
if result.best_accuracy >= orig_acc - budget:
out.append(
{
"model": model,
"accuracy": result.best_accuracy,
"method": "nas",
"architecture": result.best_architecture,
"n_parameters": result.best_n_parameters,
},
)
except Exception as exc:
log.warning("NAS failed for %s: %s", clf_name, exc)
return out