File size: 6,337 Bytes
5d97275 | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """Compact oblivious-tree ensemble (CatBoost's symmetric trees) evaluated with whole-batch numpy operations.
In an oblivious tree every node of one depth level asks the same question, so a depth-d tree is just d
(feature, border) tests and a table of 2^d leaves: the d test results, read as bits, are the leaf index.
All distinct tests of the ensemble are evaluated once per batch into a boolean matrix [n, S]; the leaf index of
every (row, tree) pair is then built with d gathers and the leaves are summed. No python loop over rows or trees,
no CatBoost at inference: the .npz file (tests + leaves, a few hundred KB) is the whole model.
"""
from __future__ import annotations
import json
import numpy as np
class Oblivious:
def __init__(self, feat, border, tree_tests, leaves, bias, names=None, classes=None, meta=None):
self.feat = np.asarray(feat, np.int32) # [S] feature of test s (last test: always false)
self.border = np.asarray(border, np.float32) # [S] test s is x[feat] > border
self.tree_tests = np.asarray(tree_tests, np.int32) # [T, D] tests of each tree, bit d = depth d
self.leaves = np.asarray(leaves, np.float32) # [T, 2^D, C]
self.bias = np.asarray(bias, np.float32) # [C]
self.names, self.classes, self.meta = list(names or []), list(classes or []), dict(meta or {})
T, L, C = self.leaves.shape
D = self.tree_tests.shape[1]
self._flat = np.ascontiguousarray(self.leaves.transpose(2, 0, 1).reshape(C, T * L)) # [C, T*L]
self._off = (np.arange(T, dtype=np.int32) * L)[:, None]
# tests sorted by feature: each feature's tests are one broadcast compare into a slice of the test matrix
order = np.lexsort((self.border, self.feat))
inv = np.empty_like(order)
inv[order] = np.arange(len(order))
fs, bs = self.feat[order], self.border[order]
self._bcol, self._tt = bs[:, None], inv[self.tree_tests]
self._groups = [(int(f), int(np.searchsorted(fs, f, "left")), int(np.searchsorted(fs, f, "right")))
for f in np.unique(fs)]
self._w = (1 << np.arange(D)).astype(np.uint8)
@property
def n_trees(self) -> int:
return self.leaves.shape[0]
def raw(self, x: np.ndarray, chunk: int = 8192) -> np.ndarray:
"""x: float32 [n, F] -> raw scores [n, C] (logits for Logloss / MultiClass).
Feature-major ([F, m]) so every gather copies whole contiguous rows; leaf indices built as uint8."""
x = np.asarray(x, np.float32)
out = np.empty((len(x), self.leaves.shape[2]), np.float32)
for s in range(0, len(x), chunk):
xt = np.ascontiguousarray(x[s:s + chunk].T) # [F, m]
b = np.empty((len(self._bcol), xt.shape[1]), bool) # [S, m] test results
for f, lo, hi in self._groups:
np.greater(xt[f], self._bcol[lo:hi], out=b[lo:hi])
idx = np.einsum("tdm,d->tm", b.view(np.uint8)[self._tt], self._w) # [T, m] leaf index, bit d = depth d
flat = idx.astype(np.int32)
flat += self._off
for c in range(out.shape[1]):
out[s:s + chunk, c] = self._flat[c][flat].sum(0)
return out + self.bias
def proba_attack(self, x: np.ndarray) -> np.ndarray:
"""P(attack) for a binary detector (one logit) or 1 - P(benign) for a multi-class model."""
r = self.raw(x)
if r.shape[1] == 1:
return 1.0 / (1.0 + np.exp(-r[:, 0]))
z = np.exp(r - r.max(1, keepdims=True))
p = z / z.sum(1, keepdims=True)
return 1.0 - p[:, self.classes.index("benign")] if "benign" in self.classes else p.max(1)
# --- io ---
def save(self, path: str, leaf_dtype=np.float32) -> None:
np.savez_compressed(path, feat=self.feat.astype(np.int16), border=self.border, tree_tests=self.tree_tests.astype(np.int16),
leaves=self.leaves.astype(leaf_dtype), bias=self.bias,
info=np.frombuffer(json.dumps({"names": self.names, "classes": self.classes,
"meta": self.meta}).encode(), np.uint8))
@classmethod
def load(cls, path: str) -> "Oblivious":
z = np.load(path)
info = json.loads(bytes(z["info"]).decode())
return cls(z["feat"], z["border"], z["tree_tests"], z["leaves"].astype(np.float32), z["bias"],
info["names"], info["classes"], info["meta"])
@classmethod
def from_catboost(cls, model_json: dict, names=None, classes=None, meta=None) -> "Oblivious":
"""From CatBoost's model.save_model(path, format='json') (float features only, symmetric trees)."""
ff = model_json["features_info"]["float_features"]
flat = {f["feature_index"]: f["flat_feature_index"] for f in ff}
trees = model_json["oblivious_trees"]
for t in trees:
t["splits"] = t.get("splits") or [] # a depth-0 tree (one leaf) is written with splits: null
D = max(1, max(len(t["splits"]) for t in trees))
C = len(trees[0]["leaf_values"]) // (1 << len(trees[0]["splits"]))
tests: dict = {}
tt = np.full((len(trees), D), -1, np.int64)
leaves = np.zeros((len(trees), 1 << D, C), np.float32)
for i, t in enumerate(trees):
for d, sp in enumerate(t["splits"]):
if sp.get("split_type", "FloatFeature") != "FloatFeature":
raise ValueError(f"unsupported split {sp}")
key = (flat[sp["float_feature_index"]], float(sp["border"]))
tt[i, d] = tests.setdefault(key, len(tests))
n = 1 << len(t["splits"])
leaves[i, :n] = np.asarray(t["leaf_values"], np.float32).reshape(n, C) # leaf-major, class inner
never = len(tests) # padding test for shallower trees: x[0] > +inf is always false
tt[tt < 0] = never
feat = [k[0] for k in tests] + [0]
border = [k[1] for k in tests] + [np.inf]
scale, bias = model_json.get("scale_and_bias", [1.0, [0.0] * C])
bias = np.broadcast_to(np.asarray(bias, np.float32), (C,))
return cls(feat, border, tt, leaves * np.float32(scale), bias, names, classes, meta)
|