33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
from typing import List, Tuple
|
|
from tinygrad.codegen.kernel import Kernel
|
|
from tinygrad.engine.search import get_kernel_actions, actions
|
|
|
|
_net = None
|
|
def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:
|
|
global _net
|
|
if _net is None:
|
|
from tinygrad.nn.state import load_state_dict, safe_load
|
|
from extra.optimization.pretrain_valuenet import ValueNet
|
|
_net = ValueNet(1021+len(actions), 2)
|
|
load_state_dict(_net, safe_load("/tmp/qnet.safetensors"), verbose=False)
|
|
from tinygrad.tensor import Tensor
|
|
from tinygrad.helpers import Context
|
|
from extra.optimization.helpers import lin_to_feats
|
|
import numpy as np
|
|
feats = []
|
|
lins = []
|
|
base_tms = []
|
|
for lin,tm in beam:
|
|
lin_feats = lin_to_feats(lin)
|
|
for a,v in get_kernel_actions(lin, include_0=False).items():
|
|
acts = np.zeros(len(actions))
|
|
acts[a-1] = 1.0
|
|
feats.append(np.concatenate([lin_feats, acts]))
|
|
lins.append(v)
|
|
base_tms.append(tm)
|
|
with Context(BEAM=0):
|
|
with Tensor.train(False):
|
|
preds = _net(Tensor(feats)).numpy()
|
|
pred_time = np.array(base_tms) / np.exp(preds[:, 0])
|
|
return sorted(zip(lins, pred_time), key=lambda x: x[1])
|