lczerolens.model#

Class for wrapping the LCZero models.

Classes#

`LczeroModel`	Class for wrapping the LCZero models.
`Flow`	Base class for isolating a flow.
`PolicyFlow`	Class for isolating the policy flow.
`ValueFlow`	Class for isolating the value flow.
`WdlFlow`	Class for isolating the WDL flow.
`MlhFlow`	Class for isolating the MLH flow.
`ForceValueFlow`	Class for forcing and isolating the value flow.

Module Contents#

class lczerolens.model.LczeroModel(model_key, *args, dispatch=False, meta_buffers=True, **kwargs)[source]#

Bases: nnsight.NNsight

Class for wrapping the LCZero models.

Parameters:

model_key (Union[str, torch.nn.Module])
dispatch (bool)
meta_buffers (bool)

trace(*inputs, **kwargs)[source]#

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)

_execute(*prepared_inputs, **kwargs)[source]#

Virtual method to run the underlying ._model with some inputs.

Default implementation util.applies moving all tensors to the device of the first parameter in ._model and passes the values into the model.

Parameters:: prepared_inputs (tuple[Any]) – Prepared inputs.
Return type:: Any

_prepare_inputs(*inputs, **kwargs)[source]#

Virtual method to prepare inputs before batching and execution and return batch size of prepared_inputs.

Default implementation just returns inputs and length of first input.

Parameters:

inputs (tuple[Any]) – Inputs to prepare for batching and execution.
int – Batch size of prepared_inputs.

Returns:

Prepared inputs, batch size of inputs.

Return type:

Tuple[tuple[Any], int]

__call__(*inputs, **kwargs)[source]#

__getattr__(key)[source]#

Wrapper of ._envoy’s attributes to access module’s inputs and outputs.

Returns:: Attribute.
Return type:: Any

__setattr__(key, value)[source]#: Overload setattr to create and set an Envoy when trying to set a torch Module.

property device[source]#: Returns the device.

classmethod from_path(model_path)[source]#

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:: LczeroModel
Raises:: NotImplementedError – If the model file extension is not supported

classmethod from_onnx_path(onnx_model_path, check=True)[source]#

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:

LczeroModel

Raises:

FileNotFoundError – If the model file does not exist
ValueError – If the model could not be loaded

classmethod from_torch_path(torch_model_path)[source]#

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:

LczeroModel

Raises:

FileNotFoundError – If the model file does not exist
ValueError – If the model could not be loaded or is not a valid model type

_ensure_proper_forward()[source]#

class lczerolens.model.Flow(model_key, *args, **kwargs)[source]#

Bases: LczeroModel

Base class for isolating a flow.

_flow_type: str[source]#

_registry: Dict[str, Type[Flow]][source]#

classmethod register(name)[source]#

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.

classmethod from_name(name, *args, **kwargs)[source]#

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:

Flow

Raises:

KeyError – If the flow name is not found.

classmethod from_model(name, model, *args, **kwargs)[source]#

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:

Flow

Raises:

KeyError – If the flow name is not found.

classmethod is_compatible(model)[source]#

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

_ensure_proper_forward()[source]#: Rewrites the forward function to return the flow output.

class lczerolens.model.PolicyFlow(model_key, *args, **kwargs)[source]#

Bases: Flow

Class for isolating the policy flow.

class lczerolens.model.ValueFlow(model_key, *args, **kwargs)[source]#

Bases: Flow

Class for isolating the value flow.

class lczerolens.model.WdlFlow(model_key, *args, **kwargs)[source]#

Bases: Flow

Class for isolating the WDL flow.

class lczerolens.model.MlhFlow(model_key, *args, **kwargs)[source]#

Bases: Flow

Class for isolating the MLH flow.

class lczerolens.model.ForceValueFlow(model_key, *args, **kwargs)[source]#

Bases: Flow

Class for forcing and isolating the value flow.

classmethod is_compatible(model)[source]#

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

_ensure_proper_forward()[source]#: Rewrites the forward function to return the flow output.