class RecDistillExperimentRunner:
def __init__(self, args: Any, wandb_logger=None):
self.config = args if isinstance(args, RecDistillConfig) else None
self.args = runner_args_from_config(args) if isinstance(args, RecDistillConfig) else args
self.wandb_logger = wandb_logger
args = self.args
self.device = torch.device(args.device) if args.device else torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.student_backbone = normalize_backbone_name(args.student_backbone)
self.teacher_source = _teacher_source_from_args(args)
self.teacher_path = self.teacher_source.path
self.output_path = _distilled_student_path(
resolve_student_checkpoint_from_args(args, distiller_name=self.resolve_distiller_name())
)
self.output_path.parent.mkdir(parents=True, exist_ok=True)
self.run_dir = self.output_path.parent.parent if self.output_path.parent.name == "artifacts" else self.output_path.parent
(self.run_dir / "perf").mkdir(parents=True, exist_ok=True)
(self.run_dir / "logs").mkdir(parents=True, exist_ok=True)
(self.run_dir / "config").mkdir(parents=True, exist_ok=True)
self.teacher_state = None
self.dataset = None
self.val_dict: dict[int, set[int]] = {}
self.test_dict: dict[int, set[int]] = {}
self.model = None
self.distiller = None
self.optimizer = None
self.trainer = None
@classmethod
def from_config(cls, config: RecDistillConfig, wandb_logger=None) -> "RecDistillExperimentRunner":
return cls(config, wandb_logger=wandb_logger)
def resolve_distiller_name(self) -> str:
names = []
if float(getattr(self.args, "lambda_de", 0.0)) > 0:
names.append("DE")
if float(getattr(self.args, "lambda_rrd", 0.0)) > 0:
names.append("RRD")
if float(getattr(self.args, "lambda_unkd", 0.0)) > 0:
names.append("UnKD")
if float(getattr(self.args, "lambda_td", 0.0)) > 0:
names.append(str(getattr(self.args, "td_type", "TD")).upper())
return "-".join(names) if names else "NONE"
def run_config(self) -> dict[str, Any]:
return {
"config_source": "RecDistillConfig" if self.config is not None else "args",
"dataset": self.args.dataset,
"teacher_model": self.args.teacher_model,
"teacher_path": str(self.teacher_path) if self.teacher_path is not None else None,
"teacher_framework": getattr(self.args, "teacher_framework", "auto"),
"teacher_format": getattr(self.args, "teacher_format", "auto"),
"teacher_embedding_dim": self.args.teacher_embedding_dim,
"student_backbone": self.student_backbone,
"student_framework": self.args.student_framework,
"student_embedding_dim": self.args.student_embedding_dim,
"lightgcn_layers": self.args.lightgcn_layers,
"neumf_mlp_dims": self.args.neumf_mlp_dims,
"neumf_dropout": self.args.neumf_dropout,
"epochs": self.args.epochs,
"batch_size": self.args.batch_size,
"learning_rate": self.args.learning_rate,
"l2_reg": self.args.l2_reg,
"lambda_de": self.args.lambda_de,
"num_experts": self.args.num_experts,
"temperature": self.args.temperature,
"lambda_rrd": self.args.lambda_rrd,
"rrd_interesting_size": self.args.rrd_interesting_size,
"rrd_uninteresting_size": self.args.rrd_uninteresting_size,
"rrd_temperature": self.args.rrd_temperature,
"rrd_teacher_topk": self.args.rrd_teacher_topk,
"lambda_unkd": self.args.lambda_unkd,
"unkd_sample_num": self.args.unkd_sample_num,
"unkd_group_count": self.args.unkd_group_count,
"unkd_popularity_lambda": self.args.unkd_popularity_lambda,
"unkd_rank_top_k": self.args.unkd_rank_top_k,
"unkd_rank_temperature": self.args.unkd_rank_temperature,
"lambda_td": self.args.lambda_td,
"td_type": self.args.td_type,
"htd_alpha": self.args.htd_alpha,
"htd_num_groups": self.args.htd_num_groups,
"htd_topology_mode": self.args.htd_topology_mode,
"htd_initial_tau": self.args.htd_initial_tau,
"htd_min_tau": self.args.htd_min_tau,
"htd_decay_epochs": self.args.htd_decay_epochs,
"td_entity_sample_size": self.args.td_entity_sample_size,
"seed": self.args.seed,
"eval_enabled": not self.args.skip_eval,
"eval_k": self.args.eval_k,
"eval_every": self.args.eval_every,
"selection_split": self.args.selection_split,
"selection_metric": self.args.selection_metric,
"assert_no_train_leak": self.args.assert_no_train_leak,
}
def prepare(self) -> None:
print("\n" + "=" * 80)
print("Student Distillation Training")
print("=" * 80)
print(f"Dataset: {self.args.dataset}")
print(f"Teacher source: {self.teacher_path or self.teacher_source.format}")
print(f"Student backbone: {self.student_backbone}")
print(f"Student framework: {self.args.student_framework}")
print(f"Device: {self.device}")
print("=" * 80 + "\n")
teacher_state = load_teacher(self.teacher_source, device="cpu")
validate_loaded_teacher_for_distillation(teacher_state, self.resolve_distiller_name())
print(f"Teacher users/items: {teacher_state.num_users}/{teacher_state.num_items}")
print(f"Teacher embedding dim: {teacher_state.embedding_dim if teacher_state.has_embeddings else 'none'}")
print(f"Teacher exact scorer: {teacher_state.scorer is not None}")
teacher_embedding_source = teacher_state.metadata.get("embedding_source")
if teacher_embedding_source:
print(f"Teacher embedding source: {teacher_embedding_source}")
teacher_representation = teacher_state.metadata.get("representation")
if teacher_representation:
print(f"Teacher embedding representation: {teacher_representation}")
self.teacher_state = teacher_state
user_mapping, item_mapping, mapping_source = resolve_teacher_dataset_mappings(
teacher_state.metadata,
dataset_name=self.args.dataset,
)
split_id_space = "dataset_integer" if mapping_source == "dataset_integer" else None
print(f"Dataset mapping source: {mapping_source}")
self.dataset, dropped = load_train_dataset(
dataset_name=self.args.dataset,
teacher_num_users=teacher_state.num_users,
teacher_num_items=teacher_state.num_items,
user_mapping=user_mapping,
item_mapping=item_mapping,
id_space=split_id_space,
)
print(f"Train interactions: {len(self.dataset.interactions)}")
print(f"Dropped interactions (out of teacher range): {dropped}")
if not self.args.skip_eval:
self.val_dict, dropped_val = load_eval_split(
dataset_name=self.args.dataset,
split_name="val",
teacher_num_users=teacher_state.num_users,
teacher_num_items=teacher_state.num_items,
user_mapping=user_mapping,
item_mapping=item_mapping,
id_space=split_id_space,
)
self.test_dict, dropped_test = load_eval_split(
dataset_name=self.args.dataset,
split_name="test",
teacher_num_users=teacher_state.num_users,
teacher_num_items=teacher_state.num_items,
user_mapping=user_mapping,
item_mapping=item_mapping,
id_space=split_id_space,
)
print(f"Validation interactions: {sum(len(v) for v in self.val_dict.values())} (dropped: {dropped_val})")
print(f"Test interactions: {sum(len(v) for v in self.test_dict.values())} (dropped: {dropped_test})")
train_loader = build_train_loader(self.dataset, batch_size=self.args.batch_size, num_workers=self.args.num_workers)
self.model = build_student_model(
backbone=self.student_backbone,
dataset=self.dataset,
embedding_dim=self.args.student_embedding_dim,
l2_reg=self.args.l2_reg,
lightgcn_layers=self.args.lightgcn_layers,
neumf_mlp_dims=self.args.neumf_mlp_dims,
neumf_dropout=self.args.neumf_dropout,
framework=self.args.student_framework,
graph_builder=build_lightgcn_graph,
).to(self.device)
self.distiller = build_distiller_from_args(
args=self.args,
teacher_state=teacher_state,
student_dim=self.args.student_embedding_dim,
)
teacher_state_for_distiller = teacher_state.to(self.device)
if self.distiller is not None:
self.distiller = self.distiller.to(self.device)
self.distiller.on_train_start(teacher_state_for_distiller, self.dataset)
prepare_distiller_trainable_modules(self.distiller, int(self.args.student_embedding_dim), self.device)
setattr(self.distiller, "_recdistill_initialized", True)
trainable_params = list(self.model.parameters())
if self.distiller is not None:
trainable_params += list(self.distiller.parameters())
self.optimizer = torch.optim.Adam(trainable_params, lr=self.args.learning_rate)
self.trainer = DistillationTrainer(
model=self.model,
optimizer=self.optimizer,
train_loader=train_loader,
distiller=self.distiller,
device=self.device,
teacher_state=teacher_state,
dataset=self.dataset,
)
if not self.args.skip_eval:
scorer_note = "exact scorer" if teacher_state.scorer is not None else "embedding dot product"
teacher_eval = evaluate_embeddings(
user_embeddings=teacher_state.user_embeddings,
item_embeddings=teacher_state.item_embeddings,
train_seen=self.dataset.train_dict,
ground_truth=self.val_dict,
top_k=self.args.eval_k,
batch_size=self.args.eval_batch_size,
device=self.device,
scorer=teacher_state.scorer,
)[0]
print(
f"Teacher baseline @ {self.args.eval_k} (val, {scorer_note}): "
f"P={teacher_eval['precision']:.4f} "
f"R={teacher_eval['recall']:.4f} "
f"NDCG={teacher_eval['ndcg']:.4f} "
f"HR={teacher_eval['hr']:.4f}"
)
def run(self) -> dict[str, Any]:
set_seed(int(self.args.seed))
self.prepare()
return self.train()
def train(self) -> dict[str, Any]:
args = self.args
assert self.teacher_state is not None
assert self.dataset is not None
assert self.model is not None
assert self.optimizer is not None
assert self.trainer is not None
start_payload = {
"status": "running",
"started_at_utc": utc_now_iso(),
**self.run_config(),
"teacher_embedding_dim": int(self.teacher_state.embedding_dim) if self.teacher_state.has_embeddings else None,
}
if self.wandb_logger is not None:
self.wandb_logger.log_start(start_payload)
history: list[dict[str, float | int]] = []
best_score = float("-inf")
best_epoch = 0
best_checkpoint = best_checkpoint_path(self.output_path)
saved_best_checkpoint = False
early_best_value: float | None = None
early_best_epoch = 0
early_bad_steps = 0
early_stopped = False
early_stop_reason: str | None = None
early_monitor_name = "total_loss" if args.early_stop_mode == "loss" else f"val_{args.early_stop_metric}"
early_best_checkpoint = self.output_path.with_name(f"{self.output_path.stem}.earlystop_best{DISTILLED_STUDENT_EXT}")
run_status = "completed"
run_error: str | None = None
caught_exception: Exception | None = None
final_test_eval: dict[str, Any] | None = None
try:
for epoch in range(1, args.epochs + 1):
metrics = self.trainer.train_epoch()
row = {"epoch": epoch, **metrics}
current_eval: dict[str, dict[str, float] | int] | None = None
if not args.skip_eval and args.eval_every > 0 and (epoch % args.eval_every == 0):
current_eval = evaluate_student(
model=self.model,
train_seen=self.dataset.train_dict,
val_gt=self.val_dict,
test_gt=self.test_dict,
top_k=args.eval_k,
batch_size=args.eval_batch_size,
device=self.device,
eval_val_only=args.eval_val_only,
)
self._update_eval_row(row, current_eval)
leaked_users_test = int(current_eval.get("leaked_users_test", 0))
if args.assert_no_train_leak and (current_eval["leaked_users_val"] > 0 or leaked_users_test > 0):
raise RuntimeError(
"Train-item leakage detected in recommendations. "
f"val_leaks={current_eval['leaked_users_val']} test_leaks={leaked_users_test}"
)
selected_split_metrics = current_eval[args.selection_split]
selected_score = float(selected_split_metrics[args.selection_metric])
row["selection_score"] = selected_score
if selected_score > best_score:
best_score = selected_score
best_epoch = epoch
self._save_checkpoint(best_checkpoint, epoch, history + [row], best_epoch, best_score)
saved_best_checkpoint = True
history.append(row)
if self.wandb_logger is not None:
self.wandb_logger.log_epoch(row)
self._print_epoch(epoch, metrics, current_eval)
if args.save_every > 0 and (epoch % args.save_every == 0):
periodic_path = self.output_path.with_name(f"{self.output_path.stem}.ep{epoch}{DISTILLED_STUDENT_EXT}")
self._save_checkpoint(periodic_path, epoch, history, best_epoch, best_score if best_epoch > 0 else None)
print(f"Saved periodic checkpoint: {periodic_path}")
early_state = self._maybe_early_stop(
epoch=epoch,
row=row,
current_eval=current_eval,
early_best_value=early_best_value,
early_best_epoch=early_best_epoch,
early_bad_steps=early_bad_steps,
early_monitor_name=early_monitor_name,
early_best_checkpoint=early_best_checkpoint,
)
early_best_value = early_state["best_value"]
early_best_epoch = early_state["best_epoch"]
early_bad_steps = early_state["bad_steps"]
if early_state["stopped"]:
early_stopped = True
early_stop_reason = early_state["reason"]
print(f"Early stopping triggered at epoch {epoch}: {early_stop_reason}")
break
except Exception as exc:
run_status = "failed"
run_error = str(exc)
caught_exception = exc
if (
caught_exception is None
and args.early_stop
and args.early_stop_restore_best
and early_best_epoch > 0
and early_best_checkpoint.exists()
):
payload = load_student_checkpoint(early_best_checkpoint, map_location=self.device)
self.model.load_state_dict(payload["student_state_dict"])
print(
f"Restored early-stop best checkpoint from epoch {payload['epoch']} "
f"({payload['monitor_name']}={payload['monitor_value']:.6f})"
)
if caught_exception is None and not args.skip_eval and args.eval_val_only and len(self.test_dict) > 0:
final_test_eval = evaluate_student(
model=self.model,
train_seen=self.dataset.train_dict,
val_gt=self.val_dict,
test_gt=self.test_dict,
top_k=args.eval_k,
batch_size=args.eval_batch_size,
device=self.device,
eval_val_only=False,
)
print(
f"Final Test@{args.eval_k}: "
f"P={final_test_eval['test']['precision']:.4f} "
f"R={final_test_eval['test']['recall']:.4f} "
f"NDCG={final_test_eval['test']['ndcg']:.4f} "
f"HR={final_test_eval['test']['hr']:.4f} "
f"| leaks={final_test_eval['leaked_users_test']}"
)
history_path = self.run_dir / "logs" / f"{self.output_path.stem}.history.json"
should_save_final = caught_exception is None and (
not saved_best_checkpoint or (args.early_stop and args.early_stop_restore_best and early_best_epoch > 0)
)
if should_save_final:
final_epoch = int(history[-1]["epoch"]) if history else 0
self._save_checkpoint(
self.output_path,
final_epoch,
history,
best_epoch,
best_score if best_epoch > 0 else None,
extra={
"early_stopped": early_stopped,
"early_stop_reason": early_stop_reason,
"early_best_epoch": early_best_epoch if early_best_epoch > 0 else None,
"early_best_value": early_best_value,
"early_monitor_name": early_monitor_name if args.early_stop else None,
"final_test_eval": final_test_eval,
},
)
if caught_exception is None:
history_path.write_text(json.dumps(history, indent=2), encoding="utf-8")
end_payload = {
"status": run_status,
"ended_at_utc": utc_now_iso(),
"best_epoch": int(best_epoch),
"best_selection_score": float(best_score) if best_epoch > 0 else None,
"best_checkpoint": str(best_checkpoint) if best_epoch > 0 else None,
"final_checkpoint": str(self.output_path) if caught_exception is None else None,
"history_file": str(history_path) if history_path.exists() else None,
"early_stopped": early_stopped,
"early_stop_reason": early_stop_reason,
"early_best_epoch": early_best_epoch if early_best_epoch > 0 else None,
"early_best_value": early_best_value,
"early_monitor_name": early_monitor_name if args.early_stop else None,
"final_test_eval": final_test_eval,
"error": run_error,
}
if self.wandb_logger is not None:
self.wandb_logger.log_end(end_payload)
if caught_exception is not None:
raise caught_exception
print("\nTraining complete.")
print(f"Student checkpoint: {self.output_path}")
if best_epoch > 0:
print(
f"Best checkpoint ({args.selection_split}.{args.selection_metric}): "
f"epoch={best_epoch} score={best_score:.6f} path={best_checkpoint}"
)
if final_test_eval is not None:
print(
f"Final test metrics: "
f"NDCG={final_test_eval['test']['ndcg']:.4f} "
f"HR={final_test_eval['test']['hr']:.4f}"
)
print(f"History JSON: {history_path}\n")
return end_payload
def _checkpoint_payload(
self,
epoch: int,
history: list[dict[str, Any]],
best_epoch: int,
best_score: float | None,
extra: dict[str, Any] | None = None,
) -> dict[str, Any]:
payload = {
"epoch": epoch,
"student_state_dict": self.model.state_dict(),
"optimizer_state_dict": self.optimizer.state_dict(),
"history": history,
"config": vars(self.args),
"teacher_path": str(self.teacher_path) if self.teacher_path is not None else None,
"teacher_dim": self.teacher_state.embedding_dim if self.teacher_state.has_embeddings else None,
"num_users": self.dataset.num_users,
"num_items": self.dataset.num_items,
"best_epoch": best_epoch,
"best_selection_score": best_score,
"best_selection_split": self.args.selection_split,
"best_selection_metric": self.args.selection_metric,
}
if extra:
payload.update(extra)
return payload
def _save_checkpoint(
self,
path: Path,
epoch: int,
history: list[dict[str, Any]],
best_epoch: int,
best_score: float | None,
extra: dict[str, Any] | None = None,
) -> None:
save_student_checkpoint(path, self._checkpoint_payload(epoch, history, best_epoch, best_score, extra))
def _update_eval_row(self, row: dict[str, Any], current_eval: dict[str, Any]) -> None:
row["val_precision"] = float(current_eval["val"]["precision"])
row["val_recall"] = float(current_eval["val"]["recall"])
row["val_ndcg"] = float(current_eval["val"]["ndcg"])
row["val_hr"] = float(current_eval["val"]["hr"])
row["leaked_users_val"] = int(current_eval["leaked_users_val"])
if not self.args.eval_val_only:
row["test_precision"] = float(current_eval["test"]["precision"])
row["test_recall"] = float(current_eval["test"]["recall"])
row["test_ndcg"] = float(current_eval["test"]["ndcg"])
row["test_hr"] = float(current_eval["test"]["hr"])
row["leaked_users_test"] = int(current_eval["leaked_users_test"])
def _print_epoch(self, epoch: int, metrics: dict[str, float], current_eval: dict[str, Any] | None) -> None:
print(
f"Epoch {epoch:03d}/{self.args.epochs:03d} | "
f"base={metrics['base_loss']:.6f} "
f"distill={metrics['distill_loss']:.6f} "
f"total={metrics['total_loss']:.6f}"
)
if current_eval is None:
return
print(
f" Val@{self.args.eval_k}: "
f"P={current_eval['val']['precision']:.4f} "
f"R={current_eval['val']['recall']:.4f} "
f"NDCG={current_eval['val']['ndcg']:.4f} "
f"HR={current_eval['val']['hr']:.4f} "
f"| leaks={current_eval['leaked_users_val']}"
)
if not self.args.eval_val_only:
print(
f" Test@{self.args.eval_k}: "
f"P={current_eval['test']['precision']:.4f} "
f"R={current_eval['test']['recall']:.4f} "
f"NDCG={current_eval['test']['ndcg']:.4f} "
f"HR={current_eval['test']['hr']:.4f} "
f"| leaks={current_eval['leaked_users_test']}"
)
def _maybe_early_stop(
self,
*,
epoch: int,
row: dict[str, Any],
current_eval: dict[str, Any] | None,
early_best_value: float | None,
early_best_epoch: int,
early_bad_steps: int,
early_monitor_name: str,
early_best_checkpoint: Path,
) -> dict[str, Any]:
if not self.args.early_stop:
return {
"best_value": early_best_value,
"best_epoch": early_best_epoch,
"bad_steps": early_bad_steps,
"stopped": False,
"reason": None,
}
current_monitor_value: float | None = None
if self.args.early_stop_mode == "loss":
current_monitor_value = float(row["total_loss"])
elif current_eval is not None:
current_monitor_value = float(current_eval["val"][self.args.early_stop_metric])
if current_monitor_value is None:
return {
"best_value": early_best_value,
"best_epoch": early_best_epoch,
"bad_steps": early_bad_steps,
"stopped": False,
"reason": None,
}
improved = False
if early_best_value is None:
improved = True
elif self.args.early_stop_mode == "loss":
improved = current_monitor_value < (early_best_value - self.args.early_stop_min_delta)
else:
improved = current_monitor_value > (early_best_value + self.args.early_stop_min_delta)
if improved:
early_best_value = current_monitor_value
early_best_epoch = epoch
early_bad_steps = 0
save_student_checkpoint(
early_best_checkpoint,
{
"epoch": epoch,
"monitor_name": early_monitor_name,
"monitor_value": current_monitor_value,
"student_state_dict": self.model.state_dict(),
"optimizer_state_dict": self.optimizer.state_dict(),
"config": vars(self.args),
},
)
else:
early_bad_steps += 1
stopped = False
reason = None
if epoch >= self.args.early_stop_warmup and early_bad_steps >= self.args.early_stop_patience:
stopped = True
reason = (
f"no improvement on {early_monitor_name} for {early_bad_steps} step(s); "
f"best={early_best_value:.6f} at epoch={early_best_epoch}"
)
return {
"best_value": early_best_value,
"best_epoch": early_best_epoch,
"bad_steps": early_bad_steps,
"stopped": stopped,
"reason": reason,
}