"""inference modules"""
import abc
import logging
import os
from collections.abc import Mapping
from types import MappingProxyType
import numpy as np
import torch
import miniosl
feature_channels = len(miniosl.channel_id)
logger = logging.getLogger(__name__)
empty_dict = MappingProxyType({})
def p2elo(p, eps=1e-4):
return -400 * np.log10(1 / (abs(p) + eps) - 1)
def sort_moves(moves, policy):
flatten = policy.flatten().tolist()
prob = [(flatten[move.policy_move_label], move) for move in moves]
prob.sort(key=lambda e: -e[0])
return prob
def softmax(x):
b = np.max(x)
y = np.exp(x - b)
return y / y.sum()
[docs]
class InferenceModel(abc.ABC):
"""interface for inference using trained models"""
def __init__(self, device):
super().__init__()
self.device = device
[docs]
@abc.abstractmethod
def infer(self, inputs: torch.Tensor) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""return inference results for batch"""
[docs]
def infer_int8(
self, inputs: np.ndarray | torch.Tensor
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""an optimized path"""
if isinstance(inputs, np.ndarray):
if inputs.dtype != np.int8:
raise ValueError(f'expected int8, got {inputs.dtype}')
inputs = torch.from_numpy(inputs)
if inputs.dtype != torch.int8:
raise ValueError(f'expected int8, got {inputs.dtype}')
inputs = inputs.to(self.device).float()
inputs /= miniosl.One
return self.infer(inputs)
[docs]
def infer_one(self, input: np.ndarray) -> list:
"""return inference results for a single instance"""
outputs = self.infer(torch.stack((torch.from_numpy(input),)))
return [_[0] for _ in outputs]
[docs]
def eval(self, input: np.ndarray, *, take_softmax: bool = False) -> tuple:
"""return (move, value, aux) tuple, after softmax"""
out = self.infer_one(input)
moves = (
softmax(out[0]).reshape(-1, 9, 9)
if take_softmax
else out[0].reshape(-1, 9, 9)
)
value = out[1]
aux = out[2].reshape(-1, 9, 9) if len(out) >= 3 else None
return (moves, value, aux)
[docs]
class OnnxInfer(InferenceModel):
first_load = True
def __init__(self, path: str, device: str):
import re
super().__init__(device)
import onnxruntime as ort
if device == '':
device = 'cpu'
if device == 'cpu':
provider = ['CPUExecutionProvider']
elif device.startswith('cuda'):
cuda_pattern = r'^cuda:([0-9]+)$'
if match := re.match(cuda_pattern, device):
provider = [
('CUDAExecutionProvider', {'device_id': int(match.group(1))}),
'CPUExecutionProvider',
]
else:
provider = ['CUDAExecutionProvider', 'CPUExecutionProvider']
else:
provider = ort.get_available_providers()
self.ort_session = ort.InferenceSession(path, providers=provider)
if OnnxInfer.first_load:
logger.info(self.ort_session.get_providers())
OnnxInfer.first_load = False
self.binding = self.ort_session.io_binding()
logger.debug(f'{device=}')
self.device = device
[docs]
def infer(self, inputs: torch.Tensor):
return self.infer_iobinding(inputs.to(self.device))
# return self.infer_naive(inputs)
[docs]
def infer_naive(self, inputs: torch.Tensor):
"""inefficient in gpu-cpu transfer if inputs are on already gpu"""
out = self.ort_session.run(None, {'input': inputs.to('cpu').numpy()})
return out
[docs]
def infer_iobinding(self, inputs: torch.Tensor):
"""work with torch 2.5.1, onnxruntime-gpu 1.20.1"""
inputs = inputs.contiguous()
device = self.device
self.binding.bind_input(
name='input',
device_type=device,
device_id=0,
element_type=np.float32,
shape=tuple(inputs.shape),
buffer_ptr=inputs.data_ptr(),
)
move_tensor = torch.empty(
(inputs.shape[0], 2187), dtype=torch.float32, device=device
).contiguous()
self.binding.bind_output(
name='move',
device_type=device,
device_id=0,
element_type=np.float32,
shape=tuple(move_tensor.shape),
buffer_ptr=move_tensor.data_ptr(),
)
# binding.bind_output('move')
value_tensor = torch.empty(
(inputs.shape[0], 4), dtype=torch.float32, device=device
).contiguous()
self.binding.bind_output(
name='value',
device_type=device,
device_id=0,
element_type=np.float32,
shape=tuple(value_tensor.shape),
buffer_ptr=value_tensor.data_ptr(),
)
# binding.bind_output('value')
aux_tensor = torch.empty(
(inputs.shape[0], 9 * 9 * 22), dtype=torch.float32, device=device
).contiguous()
self.binding.bind_output(
name='aux',
device_type=device,
device_id=0,
element_type=np.float32,
shape=tuple(aux_tensor.shape),
buffer_ptr=aux_tensor.data_ptr(),
)
# binding.bind_output('aux')
self.ort_session.run_with_iobinding(self.binding)
return (
move_tensor.to('cpu').numpy(),
value_tensor.to('cpu').numpy(),
aux_tensor.to('cpu').numpy(),
)
# out = binding.copy_outputs_to_cpu()
[docs]
class TorchTRTInfer(InferenceModel):
def __init__(self, path: str, device: str):
import torch_tensorrt # ty: ignore[unresolved-import]
super().__init__(device)
with torch_tensorrt.logging.info(): # ty: ignore[missing-argument, invalid-context-manager]
self.trt_module = torch.jit.load(path)
self.device = device
[docs]
def infer(self, inputs: torch.Tensor):
with torch.cuda.device(torch.device(self.device)):
tensor = inputs.half().to(self.device)
outputs = self.trt_module(tensor)
ret = [_.to('cpu').numpy() for _ in outputs]
return ret
[docs]
class TorchScriptInfer(InferenceModel):
def __init__(self, path: str, device: str):
super().__init__(device)
self.ts_module = torch.jit.load(path)
[docs]
def infer(self, inputs: torch.Tensor):
with torch.no_grad():
tensor = inputs.to(self.device)
outputs = self.ts_module(tensor)
return [_.to('cpu').numpy() for _ in outputs]
[docs]
class TorchInfer(InferenceModel):
def __init__(self, model, device: str):
super().__init__(device)
self.model = model
self.model.eval()
self.device = device
[docs]
def infer(self, inputs: torch.Tensor):
with torch.no_grad():
tensor = inputs.to(self.device)
outputs = self.model(tensor)
return [_.to('cpu').numpy() for _ in outputs]
[docs]
def load(
path: str,
device: str = '',
torch_cfg: Mapping = empty_dict,
*,
compiled: bool = False,
strict: bool = True,
remove_aux_head: bool = False,
) -> InferenceModel:
"""factory method to load a model from file
:param path: filepath,
:param device: torch device such as 'cuda', 'cpu',
:param torch_cfg: network specification needed for TorchInfer.
"""
# ruff: ignore[SIM102]
if remove_aux_head:
if not path.endswith(('.pt', '.chpt')):
raise ValueError('not supported')
if path.endswith('.onnx'):
return OnnxInfer(path, device)
if path.endswith('.ts'):
if not device:
device = 'cuda'
return TorchTRTInfer(path, device)
if path.endswith('.pts'): # need to used a different extention from TRT's
return TorchScriptInfer(path, device)
if path.endswith('.pt'):
NN = (
miniosl.network.PVNetwork
if remove_aux_head
else miniosl.network.StandardNetwork
)
raw_model = NN(**torch_cfg).to(device)
if compiled:
cmodel = torch.compile(raw_model)
model = cmodel
else:
model = raw_model
saved_state = torch.load(
path, map_location=torch.device(device), weights_only=False
)
strict = strict and not remove_aux_head
model.load_state_dict(saved_state, strict=strict) # ty: ignore[unresolved-attribute]
return TorchInfer(raw_model, device)
if path.endswith('.ptd'):
model = miniosl.StandardNetwork.load_with_dict(path)
model = model.to(device)
return TorchInfer(model, device)
if path.endswith('.chpt'):
checkpoint = torch.load(
path, map_location=torch.device(device), weights_only=False
)
cfg = checkpoint['cfg']
network_cfg = cfg['network_cfg']
for obsolete_key in ['make_bottleneck']:
if obsolete_key in network_cfg:
del network_cfg[obsolete_key]
NN = (
miniosl.network.PVNetwork
if remove_aux_head
else miniosl.network.StandardNetwork
)
raw_model = NN(**network_cfg).to(device)
if compiled:
cmodel = torch.compile(raw_model)
model = cmodel
else:
model = raw_model
strict = not remove_aux_head
model.load_state_dict(checkpoint['model_state_dict'], strict=strict) # ty: ignore[unresolved-attribute]
return TorchInfer(raw_model, device)
raise ValueError('unknown filetype')
class InferenceForGameArray(miniosl.InferenceModelStub):
def __init__(self, module: InferenceModel):
super().__init__()
self.module = module
def py_infer(self, features):
features = features.reshape(-1, len(miniosl.channel_id), 9, 9)
return self.module.infer_int8(features)
def export_onnx(model, *, device, filename, remove_aux_head=False):
import onnx # to detect import error eaelier
import torch.onnx
assert onnx.__version__ is not None
model.eval()
dtype = torch.float
dummy_input = (
torch.randn(1024, feature_channels, 9, 9, device=device, dtype=dtype),
)
if not filename.endswith('.onnx'):
filename = f'{filename}.onnx'
if not remove_aux_head:
torch.onnx.export(
model,
dummy_input,
filename,
dynamic_axes={
'input': {0: 'batch_size'},
'move': {0: 'batch_size'},
'value': {0: 'batch_size'},
'aux': {0: 'batch_size'},
},
verbose=False,
input_names=['input'],
output_names=['move', 'value', 'aux'],
)
else:
torch.onnx.export(
model,
dummy_input,
filename,
dynamic_axes={
'input': {0: 'batch_size'},
'move': {0: 'batch_size'},
'value': {0: 'batch_size'},
},
verbose=False,
input_names=['input'],
output_names=['move', 'value'],
)
[docs]
def export_executorch(
model: miniosl.PVNetwork, device: str, filename: str, quiet: bool
) -> None:
"""export for executorch"""
import executorch.exir
import executorch.extension
sample_data = torch.randn(
16, feature_channels, 9, 9, device=device, dtype=torch.half
)
dynamic_shapes = {
'x': {
0: torch.export.Dim('batch', min=1, max=128),
}
}
model.eval()
model = model.half()
exported_model = torch.export.export(model, (sample_data,), dynamic_shapes=dynamic_shapes)
exported_model = exported_model.run_decompositions() # decomp_table=None
logger.info('exported')
if device == 'cpu':
from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
XnnpackPartitioner,
)
partitioner = XnnpackPartitioner()
edge_compile_config = executorch.exir.EdgeCompileConfig()
elif device.startswith('cuda'):
from executorch.backends.cuda.cuda_backend import CudaBackend
from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
from executorch.exir.backend.compile_spec_schema import CompileSpec
partitioner = CudaPartitioner([
CudaBackend.generate_method_name_compile_spec("forward"),
CompileSpec("precision", b"fp16"),
CompileSpec("device", b"cuda:0"),
# CompileSpec("use_tf32", "True".encode()),
CompileSpec("max_workspace_size", str(1<<33).encode()) # 8GB for autotuning
])
edge_compile_config = executorch.exir.EdgeCompileConfig(
_check_ir_validity=False,
_skip_dim_order=True,
)
logger.info(f'{edge_compile_config=}')
else:
raise RuntimeError(f"export not implemented {device=} {filename=}")
if not quiet:
logger.info('start edgetransform')
et_program = executorch.exir.to_edge_transform_and_lower(
exported_model,
partitioner=[partitioner],
compile_config=edge_compile_config,
)
logger.warning('got et_program')
# torch._logging.set_logs(export=logging.WARNING)
exec_program = et_program.to_executorch()
with open(filename, 'wb') as file:
exec_program.write_to_file(file)
def export_tensorrt(
model: miniosl.PVNetwork, *, device, filename: str, quiet: bool = False
) -> None:
import torch_tensorrt # ty: ignore[unresolved-import]
import torch_tensorrt.ts.logging # ty: ignore[unresolved-import]
if quiet:
torch_tensorrt.logging.set_reportable_log_level( # ty: ignore[unresolved-attribute]
torch_tensorrt.logging.Level.Error # ty: ignore[unresolved-attribute]
)
torch_tensorrt.ts.logging.set_is_colored_output_on(True)
if not device:
device = 'cuda'
elif not device.startswith('cuda'):
raise ValueError(f'unexpected device for trt {device}')
model.eval()
model = model.half()
inputs = [
torch_tensorrt.Input(
min_shape=[1, feature_channels, 9, 9],
opt_shape=[128, feature_channels, 9, 9],
max_shape=[2048, feature_channels, 9, 9],
dtype=torch.half,
)
]
savefile = filename if filename.endswith('.ts') else f'{filename}.ts'
enabled_precisions: set[torch.dtype | torch_tensorrt.dtype] = {torch.half}
with torch.cuda.device(torch.device(device)):
sample_data = torch.randn(16, feature_channels, 9, 9, device=device, dtype=torch.half)
conservative = 'MINIOSL_CONSERVATIVE_EXPORT' in os.environ
if conservative:
trt_ts_module = torch_tensorrt.compile(
torch.jit.script(model),
inputs=inputs,
enabled_precisions=enabled_precisions,
ir='ts',
# device=torch.device(device)
)
_ = trt_ts_module(sample_data)
torch.jit.save(trt_ts_module, savefile)
else:
trt_gm = torch_tensorrt.compile(
model,
inputs=inputs,
# enabled_precisions=enabled_precisions,
ir='dynamo',
)
torch_tensorrt.save(trt_gm, savefile, arg_inputs=[sample_data], output_format="torchscript")
def export_torch_script(model, *, device, filename):
model.eval()
if device == 'cuda':
inputs = torch.rand(8, feature_channels, 9, 9).to(device)
ts_module = torch.jit.trace(model, inputs)
else:
ts_module = torch.jit.optimize_for_inference(torch.jit.script(model))
if not filename.endswith('.pts'):
filename = f'{filename}.pts'
torch.jit.save(ts_module, filename)
def export_model(
model: miniosl.PVNetwork, *, device, filename: str,
quiet: bool = False, remove_aux_head: bool = False
) -> None:
if filename.endswith('.onnx'):
export_onnx(
model, device=device, filename=filename, remove_aux_head=remove_aux_head
)
elif filename.endswith('.ts'):
export_tensorrt(model, device=device, filename=filename, quiet=quiet)
elif filename.endswith('.pts'):
export_torch_script(model, device=device, filename=filename)
elif filename.endswith('.ptd'): # note: different from executorch's
model.save_with_dict(filename)
elif filename.endswith('.pte'): # executorch's
export_executorch(model, device=device, filename=filename, quiet=quiet)
else:
raise ValueError(f'unknown filetype {filename}')