"""Knowledge distillation optimizer."""
from __future__ import annotations
import logging
from typing import Any
import numpy as np
from .base import BaseOptimizer, OptimizationResult, register_optimizer
log = logging.getLogger(__name__)
def _count_params(model: Any) -> int:
if hasattr(model, "parameters"):
return int(sum(p.numel() for p in model.parameters()))
return 0
[docs]
@register_optimizer("distillation")
class DistillationOptimizer(BaseOptimizer):
"""Knowledge distillation from a teacher to smaller student models."""
def __init__(
self,
student_sizes: list[int] | None = None,
temperature: float = 3.0,
accuracy_tolerance: float = 0.01,
) -> None:
self.student_sizes = student_sizes if student_sizes is not None else [64, 32]
self.temperature = temperature
self.accuracy_tolerance = accuracy_tolerance
[docs]
def optimize(
self,
model: Any,
train_data: tuple[Any, Any],
val_data: tuple[Any, Any] | None = None,
**kwargs: Any,
) -> OptimizationResult:
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
teacher = model
train_features, train_labels = train_data
train_features = np.asarray(train_features)
train_labels = np.asarray(train_labels)
if not isinstance(teacher, nn.Module):
log.warning("Teacher is not a PyTorch module; returning as-is")
return OptimizationResult(model=teacher)
device = "cuda" if torch.cuda.is_available() else "cpu"
teacher = teacher.to(device).eval()
original_params = _count_params(teacher)
x = torch.tensor(train_features, dtype=torch.float32, device=device)
y = torch.tensor(train_labels, dtype=torch.long, device=device)
batch_size = int(kwargs.get("batch_size", 128))
epochs = int(kwargs.get("epochs", 30))
lr = float(kwargs.get("lr", 1e-3))
temperature = self.temperature
with torch.no_grad():
ordered_dl = DataLoader(TensorDataset(x, y), batch_size=batch_size, shuffle=False)
teacher_logits = []
for xb, _ in ordered_dl:
teacher_logits.append(teacher(xb))
soft = torch.cat(teacher_logits, dim=0)
teacher_acc = float((soft.argmax(1).cpu().numpy() == train_labels).mean())
target_acc = teacher_acc - self.accuracy_tolerance
input_dim = train_features.shape[1]
output_dim = int(np.max(train_labels)) + 1
student_factory = kwargs.get("student_factory")
indices = torch.arange(len(x), device=device)
indexed_dl = DataLoader(TensorDataset(x, y, indices), batch_size=batch_size, shuffle=True)
best_student = None
best_acc = -1.0
best_n_params = original_params
pareto_points: list[dict[str, Any]] = []
for hidden in self.student_sizes:
if student_factory is not None:
student = student_factory(hidden_dim=hidden)
else:
student = nn.Sequential(
nn.Linear(input_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, output_dim),
)
if not isinstance(student, nn.Module):
continue
student = student.to(device)
opt = torch.optim.Adam(student.parameters(), lr=lr)
for _ in range(epochs):
student.train()
for xb, yb, idx in indexed_dl:
opt.zero_grad()
logits = student(xb)
loss = F.kl_div(
F.log_softmax(logits / temperature, dim=1),
F.softmax(soft[idx] / temperature, dim=1),
reduction="batchmean",
) * (temperature ** 2)
loss += F.cross_entropy(logits, yb)
loss.backward()
opt.step()
student.eval()
with torch.no_grad():
preds = student(x).argmax(1).cpu().numpy()
acc = float((preds == train_labels).mean())
n_params = _count_params(student)
pareto_points.append({"accuracy": acc, "n_parameters": n_params, "hidden": hidden})
if acc >= target_acc and (best_student is None or n_params < best_n_params):
best_student = student
best_acc = acc
best_n_params = n_params
if best_student is None:
best_student = teacher
best_acc = teacher_acc
best_n_params = original_params
return OptimizationResult(
model=best_student,
original_params=original_params,
optimized_params=best_n_params,
accuracy_before=teacher_acc,
accuracy_after=best_acc,
pareto_points=pareto_points,
metadata={"temperature": self.temperature, "student_sizes": self.student_sizes},
)