lczerolens.model#
Class for wrapping the LCZero models.
- class lczerolens.model.Flow(*args, **kwargs)#
Bases:
LczeroModelBase class for isolating a flow.
- classmethod from_model(name: str, model: LczeroModel, *args, **kwargs) Flow#
Returns the flow from a model.
- Parameters:
name (str) – The name of the flow to instantiate.
model (LczeroModel) – The model to create the flow from.
*args – Positional arguments passed to flow constructor.
**kwargs – Keyword arguments passed to flow constructor.
- Returns:
The instantiated flow.
- Return type:
- Raises:
KeyError – If the flow name is not found.
- classmethod from_name(name: str, *args, **kwargs) Flow#
Returns the flow from its name.
- Parameters:
name (str) – The name of the flow to instantiate.
*args – Positional arguments passed to flow constructor.
**kwargs – Keyword arguments passed to flow constructor.
- Returns:
The instantiated flow.
- Return type:
- Raises:
KeyError – If the flow name is not found.
- classmethod is_compatible(model: Module) bool#
Checks if the model is compatible with this flow.
- Parameters:
model (nn.Module) – The model to check compatibility with.
- Returns:
Whether the model is compatible with this flow.
- Return type:
bool
- classmethod register(name: str)#
Registers the flow.
- Parameters:
name (str) – The name of the flow to register.
- Returns:
Decorator function that registers the flow subclass.
- Return type:
Callable
- Raises:
ValueError – If the flow name is already registered.
- class lczerolens.model.ForceValueFlow(*args, **kwargs)#
Bases:
FlowClass for forcing and isolating the value flow.
- classmethod is_compatible(model: Module)#
Checks if the model is compatible with this flow.
- Parameters:
model (nn.Module) – The model to check compatibility with.
- Returns:
Whether the model is compatible with this flow.
- Return type:
bool
- class lczerolens.model.LczeroModel(*args, **kwargs)#
Bases:
NNsightClass for wrapping the LCZero models.
- property device#
Returns the device.
- classmethod from_onnx_path(onnx_model_path: str, check: bool = True) LczeroModel#
Builds a model from an ONNX file path.
- Parameters:
onnx_model_path (str) – Path to the ONNX model file
check (bool, optional) – Whether to perform shape inference check, by default True
- Returns:
The wrapped model instance
- Return type:
- Raises:
FileNotFoundError – If the model file does not exist
ValueError – If the model could not be loaded
- classmethod from_path(model_path: str) LczeroModel#
Creates a wrapper from a model path.
- Parameters:
model_path (str) – Path to the model file (.onnx or .pt)
- Returns:
The wrapped model instance
- Return type:
- Raises:
NotImplementedError – If the model file extension is not supported
- classmethod from_torch_path(torch_model_path: str) LczeroModel#
Builds a model from a PyTorch file path.
- Parameters:
torch_model_path (str) – Path to the PyTorch model file
- Returns:
The wrapped model instance
- Return type:
- Raises:
FileNotFoundError – If the model file does not exist
ValueError – If the model could not be loaded or is not a valid model type
- trace(*inputs: Any, **kwargs: Dict[str, Any])#
Entrypoint into the tracing and interleaving functionality nnsight provides.
In short, allows access to the future inputs and outputs of modules in order to trace what operations you would like to perform on them. This can be as simple as accessing and saving activations for inspection, or as complicated as transforming the activations and gradients in a forward pass over multiple inputs.
- Parameters:
inputs (tuple[Any])
trace (bool, optional) – If to open a tracing context. Otherwise immediately run the model and return the raw output. Defaults to True.
invoker_args (Dict[str, Any], optional) – Keyword arguments to pass to Invoker initialization, and then downstream to the model’s .prepare_inputs(…) method. Used when giving input directly to .trace(…). Defaults to None.
kwargs (Dict[str, Any]) – Keyword arguments passed to Tracer initialization, and then downstream to the model’s ._execute(…) method.
backend (Union[Backend, str]) – Backend for this Tracer object.
remote (bool) – Use RemoteBackend with default url.
- Raises:
ValueError – If trace is False and no inputs were provided (nothing to run with)
- Returns:
Either the Tracer used for tracing, or the raw output if trace is False.
- Return type:
Union[Tracer, Any]
Examples
There are a few ways you can use
.trace(...)depending in your use case.Lets use this extremely basic model for our examples:
import torch from collections import OrderedDict input_size = 5 hidden_dims = 10 output_size = 2 model = nn.Sequential(OrderedDict([ ('layer1', torch.nn.Linear(input_size, hidden_dims)), ('sigma1', torch.nn.Sigmoid()), ('layer2', torch.nn.Linear(hidden_dims, output_size)), ('sigma2', torch.nn.Sigmoid()), ])) example_input = torch.rand((1, input_size))
The first example has us running the model with a single example input, and saving the input and output of ‘layer2’ as well as the final output using the tracing context.
from nnsight import NNsight with NNsight(model).trace(example_input) as model: l2_input = model.layer2.input.save() l2_output = model.layer2.output.save() output = model.output.save() print(l2_input) print(l2_output) print(output)
The second example allows us to divide up multiple inputs into one batch, and scope an inner invoker context to each one. We indicate this simply by not passing and positional inputs into .trace(…). The Tracer object then expects you to enter each input via Tracer.invoke(…)
example_input2 = torch.rand((1, input_size)) with NNsight(model).trace() as model: with model.invoke(example_input): output1 = model.output.save() with model.invoke(example_input2): output2 = model.output.save() print(output1) print(output2)