File size: 3,044 Bytes
975dbdd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | import os
import torch
import logging
logger = logging.getLogger(__name__)
from model import RouterNetwork
class RouterPolicy:
"""Thin wrapper around RouterNetwork for loading, inference and checkpointing."""
def __init__(self, model: RouterNetwork, device: str = "cpu"):
self.model = model.to(device)
self.device = torch.device(device)
self.model.eval()
@classmethod
def load(cls, checkpoint_path: str, embedding_dim: int | None = None, device: str = "cpu"):
"""Load RouterNetwork from checkpoint. If checkpoint contains container dict with
'model_state_dict' will use that, otherwise will try to load state_dict directly.
embedding_dim may be required if checkpoint does not include meta info.
"""
checkpoint_path = os.path.abspath(checkpoint_path)
if not os.path.exists(checkpoint_path):
raise FileNotFoundError(f"Router checkpoint not found: {checkpoint_path}")
sd = torch.load(checkpoint_path, map_location="cpu")
# Unpack common formats
if isinstance(sd, dict) and "model_state_dict" in sd:
model_state = sd["model_state_dict"]
meta = sd.get("meta", {})
else:
model_state = sd
meta = {}
# Infer embedding dim from meta if available
emb_dim = embedding_dim or meta.get("embedding_dim") or meta.get("embed_dim") or 384
model = RouterNetwork(embedding_dim=emb_dim)
missing, unexpected = model.load_state_dict(model_state, strict=False)
if missing or unexpected:
logger.info(f"RouterPolicy.load: missing keys={list(missing)} unexpected={list(unexpected)}")
return cls(model, device=device)
def predict_sync(self, emb_tensor: torch.Tensor):
"""Synchronous prediction: accepts a 1-D or 2-D embedding tensor and returns
(action_idx:int, prob:float, value:float).
"""
was_1d = False
if emb_tensor.dim() == 1:
emb_tensor = emb_tensor.unsqueeze(0)
was_1d = True
emb = emb_tensor.to(self.device)
with torch.no_grad():
logits, value = self.model(emb)
probs = torch.softmax(logits, dim=-1)
top_idx = int(probs.argmax(dim=-1)[0].item())
top_prob = float(probs[0, top_idx].item())
val = float(value.squeeze(1)[0].item()) if value is not None else 0.0
return top_idx, top_prob, val
def save_checkpoint(self, path: str, optimizer: torch.optim.Optimizer | None = None, meta: dict | None = None):
path = os.path.abspath(path)
payload = {
"model_state_dict": self.model.state_dict(),
"meta": meta or {}
}
if optimizer is not None:
try:
payload["optimizer_state_dict"] = optimizer.state_dict()
except Exception:
logger.warning("Could not include optimizer state in checkpoint")
torch.save(payload, path)
return path
|